From ad0744a95f876bda7e865b268d7df376f7243252 Mon Sep 17 00:00:00 2001
From: Asim <asim@chuhnk.me>
Date: Fri, 26 Feb 2016 00:09:06 +0000
Subject: [PATCH] make the mock sliggghtly more useful

---
 registry/mock/mock.go            | 30 ++++++++++++++++++++++++------
 selector/random_selector_test.go |  4 ++--
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/registry/mock/mock.go b/registry/mock/mock.go
index 8093ac58..ced5f308 100644
--- a/registry/mock/mock.go
+++ b/registry/mock/mock.go
@@ -4,10 +4,13 @@ import (
 	"github.com/micro/go-micro/registry"
 )
 
-type MockRegistry struct{}
+type MockRegistry struct {
+	Services map[string][]*registry.Service
+}
 
-func (m *MockRegistry) GetService(service string) ([]*registry.Service, error) {
-	return []*registry.Service{
+func (m *MockRegistry) init() {
+	// add some mock data
+	m.Services["foo"] = []*registry.Service{
 		{
 			Name:    "foo",
 			Version: "1.0.0",
@@ -46,11 +49,24 @@ func (m *MockRegistry) GetService(service string) ([]*registry.Service, error) {
 				},
 			},
 		},
-	}, nil
+	}
+}
+
+func (m *MockRegistry) GetService(service string) ([]*registry.Service, error) {
+	s, ok := m.Services[service]
+	if !ok {
+		return nil, registry.ErrNotFound
+	}
+	return s, nil
+
 }
 
 func (m *MockRegistry) ListServices() ([]*registry.Service, error) {
-	return []*registry.Service{}, nil
+	var services []*registry.Service
+	for _, service := range m.Services {
+		services = append(services, service...)
+	}
+	return services, nil
 }
 
 func (m *MockRegistry) Register(s *registry.Service, opts ...registry.RegisterOption) error {
@@ -70,5 +86,7 @@ func (m *MockRegistry) String() string {
 }
 
 func NewRegistry() *MockRegistry {
-	return &MockRegistry{}
+	m := &MockRegistry{Services: make(map[string][]*registry.Service)}
+	m.init()
+	return m
 }
diff --git a/selector/random_selector_test.go b/selector/random_selector_test.go
index 479ea750..07115936 100644
--- a/selector/random_selector_test.go
+++ b/selector/random_selector_test.go
@@ -9,13 +9,13 @@ import (
 func TestRandomSelector(t *testing.T) {
 	counts := map[string]int{}
 
-	bl := &randomSelector{
+	rs := &randomSelector{
 		so: Options{
 			Registry: mock.NewRegistry(),
 		},
 	}
 
-	next, err := bl.Select("foo")
+	next, err := rs.Select("foo")
 	if err != nil {
 		t.Errorf("Unexpected error calling random select: %v", err)
 	}