func TestWatch(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "watchTest"
assert(client.Register(serviceName, ":1111"), t)
assert(client.Register(serviceName, ":2222"), t)
set, err := client.NewServiceSet(serviceName)
assert(err, t)
updates := set.Watch(true)
assert(client.Register(serviceName, ":3333"), t)
for i := 0; i < 3; i++ {
var update *agent.ServiceUpdate
select {
case update = <-updates:
case <-time.After(3 * time.Second):
t.Fatal("Timeout exceeded", i)
}
if update.Online != true {
t.Fatal("Service update of unexected status: ", update, i)
}
if update.Name != serviceName {
t.Fatal("Service update of unexected name: ", update, i)
}
}
assert(set.Close(), t)
}
func TestUnregisterAll(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "unregisterAllTest"
assert(client.Register(serviceName, ":1111"), t)
assert(client.Register(serviceName, ":2222"), t)
assert(client.Register(serviceName, ":3333"), t)
services, err := client.Services(serviceName, 1)
assert(err, t)
if len(services) != 3 {
t.Fatal("Wrong number of services")
}
assert(client.UnregisterAll(), t)
set, err := client.NewServiceSet(serviceName)
assert(err, t)
if len(set.Services()) != 0 {
t.Fatal("There should be no services")
}
assert(set.Close(), t)
}
func TestBasicRegisterAndServiceSet(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "basicTest"
assert(client.RegisterWithAttributes(serviceName, ":1111", map[string]string{"foo": "bar"}), t)
assert(client.Register(serviceName, ":2222"), t)
set, err := client.NewServiceSet(serviceName)
assert(err, t)
waitUpdates(t, set, true, 2)()
if len(set.Services()) < 2 || len(set.Addrs()) < 2 {
t.Fatal("Registered services not online")
}
wait := waitUpdates(t, set, false, 1)
assert(client.Unregister(serviceName, ":2222"), t)
wait()
if len(set.Services()) != 1 || len(set.Addrs()) != 1 {
t.Fatal("Only 1 registered service should be left")
}
if set.Services()[0].Attrs["foo"] != "bar" {
t.Fatal("Attribute not set on service as 'bar'")
}
assert(set.Close(), t)
}
func main() {
flag.Parse()
name := flag.Arg(0)
addr := flag.Arg(1)
client, err := discoverd.NewClient()
if err != nil {
log.Fatal("Error making client:", err)
}
if err = client.Register(name, addr); err != nil {
log.Fatal("Error registering:", err)
}
log.Printf("Registered %s at %s.\n", name, addr)
exit := make(chan os.Signal, 1)
signal.Notify(exit, os.Interrupt, syscall.SIGTERM)
go func() {
<-exit
log.Println("Shutting down...")
client.Unregister(name, addr)
os.Exit(0)
}()
set, err := client.NewServiceSet(name)
if err != nil {
log.Fatal("Error getting ServiceSet:", err)
}
for range time.Tick(time.Second) {
log.Println(strings.Join(set.Addrs(), ", "))
}
}
func TestNoServices(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
set, err := client.NewServiceSet("nonexistent")
assert(err, t)
if len(set.Services()) != 0 {
t.Fatal("There should be no services")
}
assert(set.Close(), t)
}
func TestFiltering(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "filterTest"
set, err := client.NewServiceSet(serviceName)
assert(err, t)
watchSet, err := client.NewServiceSet(serviceName)
assert(err, t)
assert(client.Register(serviceName, ":1111"), t)
assert(client.RegisterWithAttributes(serviceName, ":2222", map[string]string{"foo": "qux", "id": "2"}), t)
set.Filter(map[string]string{"foo": "qux"})
waitUpdates(t, watchSet, true, 2)()
if len(set.Services()) > 1 {
t.Fatal("Filter not limiting online services in set")
}
assert(client.RegisterWithAttributes(serviceName, ":3333", map[string]string{"foo": "qux", "id": "3"}), t)
waitUpdates(t, set, true, 2)()
if s := set.Services(); len(s) < 2 {
t.Fatalf("Filter not letting new matching services in set: %#v", s[0])
}
assert(client.RegisterWithAttributes(serviceName, ":4444", map[string]string{"foo": "baz"}), t)
waitUpdates(t, watchSet, true, 4)()
if len(set.Services()) > 2 {
t.Fatal("Filter not limiting new unmatching services from set")
}
assert(set.Close(), t)
assert(watchSet.Close(), t)
}
func TestLeaderChannel(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "leadersTest"
assert(client.Register(serviceName, ":1111"), t)
set, err := client.NewServiceSet(serviceName)
assert(err, t)
leader := make(chan *discoverd.Service, 3)
go func() {
leaders := set.Leaders()
for {
leader <- <-leaders
}
}()
if (<-leader).Addr != "127.0.0.1:1111" {
t.Fatal("Incorrect leader")
}
assert(client.Unregister(serviceName, ":1111"), t)
if (<-leader) != nil {
t.Fatal("Incorrect leader")
}
assert(client.Register(serviceName, ":2222"), t)
assert(client.Register(serviceName, ":3333"), t)
if (<-leader).Addr != "127.0.0.1:2222" {
t.Fatal("Incorrect leader", leader)
}
assert(client.Unregister(serviceName, ":2222"), t)
if (<-leader).Addr != "127.0.0.1:3333" {
t.Fatal("Incorrect leader")
}
assert(set.Close(), t)
}
func TestSelecting(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "selectTest"
set, err := client.NewServiceSet(serviceName)
assert(err, t)
assert(client.Register(serviceName, ":1111"), t)
assert(client.RegisterWithAttributes(serviceName, ":2222", map[string]string{"foo": "qux", "id": "2"}), t)
assert(client.RegisterWithAttributes(serviceName, ":3333", map[string]string{"foo": "qux", "id": "3"}), t)
waitUpdates(t, set, true, 3)()
if s := set.Select(map[string]string{"id": "3"}); len(s) != 1 {
t.Fatalf("Expected one service, got: %#v", s)
}
assert(set.Close(), t)
}
func TestNewAttributes(t *testing.T) {
client, cleanup := testutil.SetupDiscoverd(t)
defer cleanup()
serviceName := "attributeTest"
set, err := client.NewServiceSet(serviceName)
assert(err, t)
assert(client.RegisterWithAttributes(serviceName, ":1111", map[string]string{"foo": "bar"}), t)
waitUpdates(t, set, true, 1)()
wait := waitUpdates(t, set, false, 1)
assert(client.RegisterWithAttributes(serviceName, ":1111", map[string]string{"foo": "baz"}), t)
wait()
if s := set.Services()[0]; s.Attrs["foo"] != "baz" {
t.Fatalf(`Expected attribute set on re-registered service to be "baz", not %q`, s.Attrs["foo"])
}
assert(set.Close(), t)
}