vendor/github.com/go-kit/kit/sd/consul/client_test.go

   1 package consul
   2
   3 import (
   4         "context"
   5         "errors"
   6         "io"
   7         "reflect"
   8         "testing"
   9
  10         stdconsul "github.com/hashicorp/consul/api"
  11
  12         "github.com/go-kit/kit/endpoint"
  13 )
  14
  15 func TestClientRegistration(t *testing.T) {
  16         c := newTestClient(nil)
  17
  18         services, _, err := c.Service(testRegistration.Name, "", true, &stdconsul.QueryOptions{})
  19         if err != nil {
  20                 t.Error(err)
  21         }
  22         if want, have := 0, len(services); want != have {
  23                 t.Errorf("want %d, have %d", want, have)
  24         }
  25
  26         if err := c.Register(testRegistration); err != nil {
  27                 t.Error(err)
  28         }
  29
  30         if err := c.Register(testRegistration); err == nil {
  31                 t.Errorf("want error, have %v", err)
  32         }
  33
  34         services, _, err = c.Service(testRegistration.Name, "", true, &stdconsul.QueryOptions{})
  35         if err != nil {
  36                 t.Error(err)
  37         }
  38         if want, have := 1, len(services); want != have {
  39                 t.Errorf("want %d, have %d", want, have)
  40         }
  41
  42         if err := c.Deregister(testRegistration); err != nil {
  43                 t.Error(err)
  44         }
  45
  46         if err := c.Deregister(testRegistration); err == nil {
  47                 t.Errorf("want error, have %v", err)
  48         }
  49
  50         services, _, err = c.Service(testRegistration.Name, "", true, &stdconsul.QueryOptions{})
  51         if err != nil {
  52                 t.Error(err)
  53         }
  54         if want, have := 0, len(services); want != have {
  55                 t.Errorf("want %d, have %d", want, have)
  56         }
  57 }
  58
  59 type testClient struct {
  60         entries []*stdconsul.ServiceEntry
  61 }
  62
  63 func newTestClient(entries []*stdconsul.ServiceEntry) *testClient {
  64         return &testClient{
  65                 entries: entries,
  66         }
  67 }
  68
  69 var _ Client = &testClient{}
  70
  71 func (c *testClient) Service(service, tag string, _ bool, opts *stdconsul.QueryOptions) ([]*stdconsul.ServiceEntry, *stdconsul.QueryMeta, error) {
  72         var results []*stdconsul.ServiceEntry
  73
  74         for _, entry := range c.entries {
  75                 if entry.Service.Service != service {
  76                         continue
  77                 }
  78                 if tag != "" {
  79                         tagMap := map[string]struct{}{}
  80
  81                         for _, t := range entry.Service.Tags {
  82                                 tagMap[t] = struct{}{}
  83                         }
  84
  85                         if _, ok := tagMap[tag]; !ok {
  86                                 continue
  87                         }
  88                 }
  89
  90                 results = append(results, entry)
  91         }
  92
  93         return results, &stdconsul.QueryMeta{}, nil
  94 }
  95
  96 func (c *testClient) Register(r *stdconsul.AgentServiceRegistration) error {
  97         toAdd := registration2entry(r)
  98
  99         for _, entry := range c.entries {
 100                 if reflect.DeepEqual(*entry, *toAdd) {
 101                         return errors.New("duplicate")
 102                 }
 103         }
 104
 105         c.entries = append(c.entries, toAdd)
 106         return nil
 107 }
 108
 109 func (c *testClient) Deregister(r *stdconsul.AgentServiceRegistration) error {
 110         toDelete := registration2entry(r)
 111
 112         var newEntries []*stdconsul.ServiceEntry
 113         for _, entry := range c.entries {
 114                 if reflect.DeepEqual(*entry, *toDelete) {
 115                         continue
 116                 }
 117                 newEntries = append(newEntries, entry)
 118         }
 119         if len(newEntries) == len(c.entries) {
 120                 return errors.New("not found")
 121         }
 122
 123         c.entries = newEntries
 124         return nil
 125 }
 126
 127 func registration2entry(r *stdconsul.AgentServiceRegistration) *stdconsul.ServiceEntry {
 128         return &stdconsul.ServiceEntry{
 129                 Node: &stdconsul.Node{
 130                         Node:    "some-node",
 131                         Address: r.Address,
 132                 },
 133                 Service: &stdconsul.AgentService{
 134                         ID:      r.ID,
 135                         Service: r.Name,
 136                         Tags:    r.Tags,
 137                         Port:    r.Port,
 138                         Address: r.Address,
 139                 },
 140                 // Checks ignored
 141         }
 142 }
 143
 144 func testFactory(instance string) (endpoint.Endpoint, io.Closer, error) {
 145         return func(context.Context, interface{}) (interface{}, error) {
 146                 return instance, nil
 147         }, nil, nil
 148 }
 149
 150 var testRegistration = &stdconsul.AgentServiceRegistration{
 151         ID:      "my-id",
 152         Name:    "my-name",
 153         Tags:    []string{"my-tag-1", "my-tag-2"},
 154         Port:    12345,
 155         Address: "my-address",
 156 }