ugly tests fix

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>

gossip.go

@@ -4,6 +4,7 @@ package gossip
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net"
@@ -27,6 +28,10 @@ import (
 // type actionType int32
 // type updateType int32
 
+var (
+	ErrNotConnected = errors.New("register not connected")
+)
+
 const (
 	actionTypeInvalid int32 = iota
 	actionTypeCreate
@@ -100,6 +105,7 @@ type gossipRegister struct {
 	interval    time.Duration
 	tcpInterval time.Duration
 
+	connected      bool
 	connectRetry   bool
 	connectTimeout time.Duration
 	sync.RWMutex
@@ -396,6 +402,7 @@ func (d *delegate) MergeRemoteState(buf []byte, join bool) {
 }
 
 func (g *gossipRegister) Connect(ctx context.Context) error {
+	g.connected = true
 	return nil
 	//return g.connect(g.opts.Addrs)
 }
@@ -491,6 +498,7 @@ func (g *gossipRegister) subscribe() (chan *register.Result, chan bool) {
 }
 
 func (g *gossipRegister) Stop() error {
+	g.connected = false
 	select {
 	case <-g.done:
 		return nil
@@ -724,6 +732,9 @@ func (g *gossipRegister) Name() string {
 }
 
 func (g *gossipRegister) Register(ctx context.Context, s *register.Service, opts ...register.RegisterOption) error {
+	if !g.connected {
+		return ErrNotConnected
+	}
 	if g.opts.Logger.V(logger.DebugLevel) {
 		g.opts.Logger.Debugf(g.opts.Context, "[gossip] Register registering service: %s", s.Name)
 	}
@@ -776,6 +787,9 @@ func (g *gossipRegister) Register(ctx context.Context, s *register.Service, opts
 }
 
 func (g *gossipRegister) Deregister(ctx context.Context, s *register.Service, opts ...register.DeregisterOption) error {
+	if !g.connected {
+		return ErrNotConnected
+	}
 
 	if g.opts.Logger.V(logger.DebugLevel) {
 		g.opts.Logger.Debugf(g.opts.Context, "[gossip] Register deregistering service: %s", s.Name)

gossip_test.go

@@ -193,15 +193,14 @@ func TestGossipRegisterRetry(t *testing.T) {
 		t.Fatalf("[gossip register] broadcast failed: service.1 not found in r2")
 	}
 
-	if err = r1.(*gossipRegister).Stop(); err != nil {
+	mu.Lock()
+	if err = r1.Disconnect(ctx); err != nil {
 		t.Fatalf("[gossip register] failed to stop register: %v", err)
 	}
-
-	mu.Lock()
 	r1 = nil
 	mu.Unlock()
 
-	<-time.After(3 * time.Second)
+	<-time.After(4 * time.Second)
 
 	found = false
 	svcs, err = r2.ListServices(ctx)
@@ -219,7 +218,15 @@ func TestGossipRegisterRetry(t *testing.T) {
 		t.Fatalf("[gossip register] service.1 found in r2")
 	}
 
+	mu.Lock()
 	r1 = newRegister(Config(mc1), Address("127.0.0.1:54321"))
+	if err := r1.Init(); err != nil {
+		t.Fatal(err)
+	}
+	if err := r1.Connect(ctx); err != nil {
+		t.Fatal(err)
+	}
+	mu.Unlock()
 	<-time.After(2 * time.Second)
 
 	found = false