Repository: open-telemetry/opamp-go
Pull request #37 (2021-12-07 20:28:48)

Fix client test case
This fixes the case when Stop returns without an error since
it manages to stop before the external context cancellation
is effective.


Patch:
diff --git a/server/serverimpl.go b/server/serverimpl.go
index 8ce004d..dc34bca 100644
--- a/server/serverimpl.go
+++ b/server/serverimpl.go
@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"errors"
+	"net"
 	"net/http"
 
 	"github.com/gorilla/websocket"
@@ -68,29 +69,54 @@ func (s *server) Start(settings StartSettings) error {
 	mux.HandleFunc(path, s.httpHandler)
 
 	hs := &http.Server{
-		Handler: mux,
-		Addr:    settings.ListenEndpoint,
+		Handler:   mux,
+		Addr:      settings.ListenEndpoint,
+		TLSConfig: settings.TLSConfig,
 	}
 	s.httpServer = hs
 
-	// Start the HTTP server in background and return immediately.
-	// TODO: wait until server starts listening or remove the requirement to wait for that.
-	if settings.TLSConfig != nil {
-		hs.TLSConfig = settings.TLSConfig
-		go s.runServer(func() error { return hs.ListenAndServeTLS("", "") })
+	listenAddr := s.httpServer.Addr
+
+	// Start the HTTP server in background.
+	if hs.TLSConfig != nil {
+		if listenAddr == "" {
+			listenAddr = ":https"
+		}
+		err = s.startHttpServer(
+			listenAddr,
+			func(l net.Listener) error { return hs.ServeTLS(l, "", "") },
+		)
 	} else {
-		go s.runServer(func() error { return hs.ListenAndServe() })
+		if listenAddr == "" {
+			listenAddr = ":http"
+		}
+		err = s.startHttpServer(
+			listenAddr,
+			func(l net.Listener) error { return hs.Serve(l) },
+		)
 	}
-	return nil
+	return err
 }
 
-func (s *server) runServer(runner func() error) {
-	err := runner()
-	// ErrServerClosed is expected after successful Stop(), so we won't log that
-	// particular error.
-	if err != nil && err != http.ErrServerClosed {
-		s.logger.Errorf("Error running HTTP Server: %v", err)
+func (s *server) startHttpServer(listenAddr string, serveFunc func(l net.Listener) error) error {
+	// If the listen address is not specified use the default.
+	ln, err := net.Listen("tcp", listenAddr)
+	if err != nil {
+		return err
 	}
+
+	// Begin serving connections in the background.
+	go func() {
+		err = serveFunc(ln)
+
+		// ErrServerClosed is expected after successful Stop(), so we won't log that
+		// particular error.
+		if err != nil && err != http.ErrServerClosed {
+			s.logger.Errorf("Error running HTTP Server: %v", err)
+		}
+	}()
+
+	return nil
 }
 
 func (s *server) Stop(ctx context.Context) error {


Test patch:
diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 9357732..89fef27 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -107,7 +107,9 @@ func TestStopCancellation(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 	err := client.Stop(ctx)
-	assert.ErrorIs(t, err, context.Canceled)
+	if err != nil {
+		assert.ErrorIs(t, err, context.Canceled)
+	}
 }
 
 func TestConnectWithServer(t *testing.T) {
diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go
index 71ae2f7..6fabec3 100644
--- a/server/serverimpl_test.go
+++ b/server/serverimpl_test.go
@@ -31,8 +31,6 @@ func startServer(t *testing.T, settings *StartSettings) *server {
 	err := srv.Start(*settings)
 	require.NoError(t, err)
 
-	// Wait until the server strts accepting connections.
-	testhelpers.WaitForEndpoint(settings.ListenEndpoint)
 	return srv
 }
 
@@ -111,6 +109,7 @@ func TestServerStartAcceptConnection(t *testing.T) {
 	// Verify that the connection is successful.
 	assert.NoError(t, err)
 	assert.NotNil(t, conn)
+	require.NotNil(t, resp)
 	assert.EqualValues(t, 101, resp.StatusCode)
 	eventually(t, func() bool { return atomic.LoadInt32(&connectedCalled) == 1 })
 	assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0)
@@ -149,6 +148,7 @@ func TestServerReceiveSendMessage(t *testing.T) {
 
 	// Connect using a WebSocket client.
 	conn, _, _ := dialClient(settings)
+	require.NotNil(t, conn)
 	defer conn.Close()
 
 	// Send a message to the server.
@@ -211,9 +211,6 @@ func TestServerAttachAcceptConnection(t *testing.T) {
 	hs := httptest.NewServer(mux)
 	defer hs.Close()
 
-	// Wait until server is ready and accepts connections.
-	testhelpers.WaitForEndpoint(hs.Listener.Addr().String())
-
 	// Connect using WebSocket client.
 	srvUrl := "ws://" + hs.Listener.Addr().String() + path
 	conn, resp, err := websocket.DefaultDialer.Dial(srvUrl, nil)