From 27540ee25553489bd8eaf24cea72c0813ce9ce53 Mon Sep 17 00:00:00 2001
From: Ivan Kozlovic <ivan@synadia.com>
Date: Fri, 3 Jul 2020 11:30:48 -0600
Subject: [PATCH] Fixed some flappers

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>
---
 server/client_test.go | 12 +++++++-----
 server/server_test.go | 16 +++++++++-------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/server/client_test.go b/server/client_test.go
index ee4eabcb..6878d452 100644
--- a/server/client_test.go
+++ b/server/client_test.go
@@ -185,12 +185,14 @@ func TestAsyncClientWithRunningServer(t *testing.T) {
 	defer c.close()
 
 	buf := make([]byte, 1000000)
-	n := runtime.Stack(buf, true)
-
 	writeLoopTxt := fmt.Sprintf("writeLoop(%p)", c.client)
-	if count := strings.Count(string(buf[:n]), writeLoopTxt); count != 1 {
-		t.Fatalf("writeLoop for client started more than once: %v", count)
-	}
+	checkFor(t, time.Second, 15*time.Millisecond, func() error {
+		n := runtime.Stack(buf, true)
+		if count := strings.Count(string(buf[:n]), writeLoopTxt); count != 1 {
+			return fmt.Errorf("writeLoop for client should have been started only once: %v", count)
+		}
+		return nil
+	})
 }
 
 func TestClientCreateAndInfo(t *testing.T) {
diff --git a/server/server_test.go b/server/server_test.go
index ed231886..bc095734 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -401,16 +401,18 @@ func TestClientAdvertiseErrorOnStartup(t *testing.T) {
 	opts.ClientAdvertise = "addr:::123"
 	s := New(opts)
 	defer s.Shutdown()
-	dl := &DummyLogger{}
-	s.SetLogger(dl, false, false)
+	l := &captureFatalLogger{fatalCh: make(chan string, 1)}
+	s.SetLogger(l, false, false)
 
 	// Expect this to return due to failure
 	s.Start()
-	dl.Lock()
-	msg := dl.msg
-	dl.Unlock()
-	if !strings.Contains(msg, "ClientAdvertise") {
-		t.Fatalf("Unexpected error: %v", msg)
+	select {
+	case msg := <-l.fatalCh:
+		if !strings.Contains(msg, "ClientAdvertise") {
+			t.Fatalf("Unexpected error: %v", msg)
+		}
+	case <-time.After(time.Second):
+		t.Fatalf("Should have failed to start")
 	}
 }