From 7c3ce60ce6ba039b1525f97e6d8f4a14550ec7da Mon Sep 17 00:00:00 2001
From: Asim <asim@chuhnk.me>
Date: Sat, 23 Apr 2016 20:57:46 +0100
Subject: [PATCH] Add label and version filters

---
 selector/filter.go      |  54 ++++++++++++++
 selector/filter_test.go | 158 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 212 insertions(+)
 create mode 100644 selector/filter.go
 create mode 100644 selector/filter_test.go

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{
+				&registry.Service{
+					Name:    "test",
+					Version: "1.0.0",
+					Nodes: []*registry.Node{
+						&registry.Node{
+							Id:      "test-1",
+							Address: "localhost",
+							Metadata: map[string]string{
+								"foo": "bar",
+							},
+						},
+					},
+				},
+				&registry.Service{
+					Name:    "test",
+					Version: "1.1.0",
+					Nodes: []*registry.Node{
+						&registry.Node{
+							Id:      "test-2",
+							Address: "localhost",
+							Metadata: map[string]string{
+								"foo": "baz",
+							},
+						},
+					},
+				},
+			},
+			label: [2]string{"foo", "bar"},
+			count: 1,
+		},
+		{
+			services: []*registry.Service{
+				&registry.Service{
+					Name:    "test",
+					Version: "1.0.0",
+					Nodes: []*registry.Node{
+						&registry.Node{
+							Id:      "test-1",
+							Address: "localhost",
+						},
+					},
+				},
+				&registry.Service{
+					Name:    "test",
+					Version: "1.1.0",
+					Nodes: []*registry.Node{
+						&registry.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{
+				&registry.Service{
+					Name:    "test",
+					Version: "1.0.0",
+				},
+				&registry.Service{
+					Name:    "test",
+					Version: "1.1.0",
+				},
+			},
+			version: "1.0.0",
+			count:   1,
+		},
+		{
+			services: []*registry.Service{
+				&registry.Service{
+					Name:    "test",
+					Version: "1.0.0",
+				},
+				&registry.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)
+		}
+	}
+}