diff --git a/selector/filter.go b/selector/filter.go new file mode 100644 index 00000000..0a644b0c --- /dev/null +++ b/selector/filter.go @@ -0,0 +1,54 @@ +package selector + +import ( + "github.com/micro/go-micro/registry" +) + +// LabelFilter is a label based Select Filter which will +// only return services with the label specified. +func LabelFilter(key, val string) Filter { + return func(old []*registry.Service) []*registry.Service { + var services []*registry.Service + + for _, service := range old { + serv := new(registry.Service) + var nodes []*registry.Node + + for _, node := range service.Nodes { + if node.Metadata == nil { + continue + } + + if node.Metadata[key] == val { + nodes = append(nodes, node) + } + } + + // only add service if there's some nodes + if len(nodes) > 0 { + // copy + *serv = *service + serv.Nodes = nodes + services = append(services, serv) + } + } + + return services + } +} + +// VersionFilter is a version based Select Filter which will +// only return services with the version specified. +func VersionFilter(version string) Filter { + return func(old []*registry.Service) []*registry.Service { + var services []*registry.Service + + for _, service := range old { + if service.Version == version { + services = append(services, service) + } + } + + return services + } +} diff --git a/selector/filter_test.go b/selector/filter_test.go new file mode 100644 index 00000000..a77ca645 --- /dev/null +++ b/selector/filter_test.go @@ -0,0 +1,158 @@ +package selector + +import ( + "testing" + + "github.com/micro/go-micro/registry" +) + +func TestLabelFilter(t *testing.T) { + testData := []struct { + services []*registry.Service + label [2]string + count int + }{ + { + services: []*registry.Service{ + ®istry.Service{ + Name: "test", + Version: "1.0.0", + Nodes: []*registry.Node{ + ®istry.Node{ + Id: "test-1", + Address: "localhost", + Metadata: map[string]string{ + "foo": "bar", + }, + }, + }, + }, + ®istry.Service{ + Name: "test", + Version: "1.1.0", + Nodes: []*registry.Node{ + ®istry.Node{ + Id: "test-2", + Address: "localhost", + Metadata: map[string]string{ + "foo": "baz", + }, + }, + }, + }, + }, + label: [2]string{"foo", "bar"}, + count: 1, + }, + { + services: []*registry.Service{ + ®istry.Service{ + Name: "test", + Version: "1.0.0", + Nodes: []*registry.Node{ + ®istry.Node{ + Id: "test-1", + Address: "localhost", + }, + }, + }, + ®istry.Service{ + Name: "test", + Version: "1.1.0", + Nodes: []*registry.Node{ + ®istry.Node{ + Id: "test-2", + Address: "localhost", + }, + }, + }, + }, + label: [2]string{"foo", "bar"}, + count: 0, + }, + } + + for _, data := range testData { + filter := LabelFilter(data.label[0], data.label[1]) + services := filter(data.services) + + if len(services) != data.count { + t.Fatalf("Expected %d services, got %d", data.count, len(services)) + } + + for _, service := range services { + var seen bool + + for _, node := range service.Nodes { + if node.Metadata[data.label[0]] != data.label[1] { + t.Fatal("Expected %s=%s but got %s=%s for service %+v node %+v", + data.label[0], data.label[1], data.label[0], node.Metadata[data.label[0]], service, node) + } + seen = true + } + + if !seen { + t.Fatalf("Expected node for %s=%s but saw none; results %+v", data.label[0], data.label[1], service) + } + } + } +} + +func TestVersionFilter(t *testing.T) { + testData := []struct { + services []*registry.Service + version string + count int + }{ + { + services: []*registry.Service{ + ®istry.Service{ + Name: "test", + Version: "1.0.0", + }, + ®istry.Service{ + Name: "test", + Version: "1.1.0", + }, + }, + version: "1.0.0", + count: 1, + }, + { + services: []*registry.Service{ + ®istry.Service{ + Name: "test", + Version: "1.0.0", + }, + ®istry.Service{ + Name: "test", + Version: "1.1.0", + }, + }, + version: "2.0.0", + count: 0, + }, + } + + for _, data := range testData { + filter := VersionFilter(data.version) + services := filter(data.services) + + if len(services) != data.count { + t.Fatalf("Expected %d services, got %d", data.count, len(services)) + } + + var seen bool + + for _, service := range services { + if service.Version != data.version { + t.Fatalf("Expected version %s, got %s", data.version, service.Version) + } + seen = true + } + + if seen == false && data.count > 0 { + t.Fatalf("Expected %d services but seen is %t; result %+v", data.count, seen, services) + } + } +}