test: Use lame duck and wait shutdown on JS tests

Signed-off-by: Waldemar Quevedo <wally@synadia.com>

server/jetstream_cluster_test.go

@@ -5928,7 +5928,7 @@ func createJetStreamSuperCluster(t *testing.T, numServersPer, numClusters int) *
 			bconf := fmt.Sprintf(jsClusterTempl, sn, storeDir, cn, cp+si, routeConfig)
 			conf := fmt.Sprintf(jsSuperClusterTempl, bconf, cn, gp, gwconf)
 			gp++
-			s, o := RunServerWithConfig(createConfFile(t, []byte(conf)))
+			s, o := RunServerWithLameDuckConfig(createConfFile(t, []byte(conf)))
 			c.servers = append(c.servers, s)
 			c.opts = append(c.opts, o)
 		}
@@ -6265,7 +6265,7 @@ var skip = func(t *testing.T) {
 	t.SkipNow()
 }
 
-func jsClientConnect(t *testing.T, s *Server, opts ...nats.Option) (*nats.Conn, nats.JetStreamContext) {
+func jsClientConnect(t *testing.T, s testJSServer, opts ...nats.Option) (*nats.Conn, nats.JetStreamContext) {
 	t.Helper()
 	nc, err := nats.Connect(s.ClientURL(), opts...)
 	if err != nil {

server/jetstream_test.go

@@ -79,24 +79,61 @@ func TestJetStreamBasicNilConfig(t *testing.T) {
 	}
 }
 
-func RunBasicJetStreamServer() *Server {
+type testJSServer interface {
+	ClientURL() string
+	Shutdown()
+}
+
+// jsServer is a specialized version of the server used for JetStreamt tests.
+type jsServer struct {
+	*Server
+	myopts  *Options
+	restart sync.Mutex
+}
+
+// Restart can be used to start again a server
+// using the same listen address as before.
+func (srv *jsServer) Restart() {
+	srv.restart.Lock()
+	defer srv.restart.Unlock()
+	srv.Server = RunServer(srv.myopts)
+}
+
+// Shutdown will graceful stop a NATS Server with JetStream.
+func (srv *jsServer) Shutdown() {
+	srv.Server.Shutdown()
+	srv.Server.WaitForShutdown()
+}
+
+// ClientURL is used to get the address to connect to a NATS Server with JetStream.
+func (srv *jsServer) ClientURL() string {
+	return srv.Server.ClientURL()
+}
+
+func RunBasicJetStreamServer() *jsServer {
 	opts := DefaultTestOptions
 	opts.Port = -1
 	opts.JetStream = true
+	opts.LameDuckDuration = 3 * time.Second
+	opts.LameDuckGracePeriod = 2 * time.Second
 	tdir, _ := ioutil.TempDir(tempRoot, "jstests-storedir-")
 	opts.StoreDir = tdir
-	return RunServer(&opts)
+	srv := RunServer(&opts)
+	return &jsServer{Server: srv, myopts: &opts}
 }
 
-func RunJetStreamServerOnPort(port int, sd string) *Server {
+func RunJetStreamServerOnPort(port int, sd string) *jsServer {
 	opts := DefaultTestOptions
 	opts.Port = port
 	opts.JetStream = true
+	opts.LameDuckDuration = 3 * time.Second
+	opts.LameDuckGracePeriod = 2 * time.Second
 	opts.StoreDir = filepath.Dir(sd)
-	return RunServer(&opts)
+	srv := RunServer(&opts)
+	return &jsServer{Server: srv, myopts: &opts}
 }
 
-func clientConnectToServer(t *testing.T, s *Server) *nats.Conn {
+func clientConnectToServer(t *testing.T, s testJSServer) *nats.Conn {
 	t.Helper()
 	nc, err := nats.Connect(s.ClientURL(),
 		nats.Name("JS-TEST"),
@@ -108,7 +145,7 @@ func clientConnectToServer(t *testing.T, s *Server) *nats.Conn {
 	return nc
 }
 
-func clientConnectWithOldRequest(t *testing.T, s *Server) *nats.Conn {
+func clientConnectWithOldRequest(t *testing.T, s testJSServer) *nats.Conn {
 	nc, err := nats.Connect(s.ClientURL(), nats.UseOldRequestStyle())
 	if err != nil {
 		t.Fatalf("Failed to create client: %v", err)

server/server_test.go

@@ -107,6 +107,14 @@ func RunServerWithConfig(configFile string) (srv *Server, opts *Options) {
 	return
 }
 
+func RunServerWithLameDuckConfig(configFile string) (srv *Server, opts *Options) {
+	opts = LoadConfig(configFile)
+	opts.LameDuckDuration = 3 * time.Second
+	opts.LameDuckGracePeriod = 2 * time.Second
+	srv = RunServer(opts)
+	return
+}
+
 func TestVersionMatchesTag(t *testing.T) {
 	tag := os.Getenv("TRAVIS_TAG")
 	// Travis started to return '' when no tag is set. Support both now.

server/test_test.go

@@ -189,6 +189,7 @@ func (c *cluster) shutdown() {
 			os.RemoveAll(sd)
 		}
 		s.Shutdown()
+		s.WaitForShutdown()
 	}
 }