From b120a03f43dfa0cc1b6d7ff407814b77593300a1 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Wed, 18 Sep 2024 11:09:15 -0400
Subject: [PATCH] Fix tests (#5114)

---
 internal/clientcache/cmd/search/search_test.go       | 10 +++++++---
 internal/clientcache/internal/daemon/options.go      | 12 ++++++++++++
 internal/clientcache/internal/daemon/options_test.go | 10 ++++++++++
 internal/clientcache/internal/daemon/server.go       |  3 +++
 4 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/internal/clientcache/cmd/search/search_test.go b/internal/clientcache/cmd/search/search_test.go
index 45ad346b5e..d36ea5fa64 100644
--- a/internal/clientcache/cmd/search/search_test.go
+++ b/internal/clientcache/cmd/search/search_test.go
@@ -61,15 +61,19 @@ func TestSearch(t *testing.T) {
 		return nil, errors.New("test not found error")
 	}
 
+	readyNotificationCh := make(chan struct{})
 	srv := daemon.NewTestServer(t, cmd)
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		srv.Serve(t, daemon.WithBoundaryTokenReaderFunc(ctx, boundaryTokenReaderFn))
+		srv.Serve(
+			t,
+			daemon.WithBoundaryTokenReaderFunc(ctx, boundaryTokenReaderFn),
+			daemon.WithReadyToServeNotificationCh(context.Background(), readyNotificationCh),
+		)
 	}()
-	// Give the store some time to get initialized
-	time.Sleep(100 * time.Millisecond)
+	<-readyNotificationCh
 	srv.AddKeyringToken(t, "address", "keyringtype", "tokenname", at.Id, boundaryTokenReaderFn)
 	srv.AddKeyringToken(t, "address", "keyringtype", "unsupported", unsupportedAt.Id, boundaryTokenReaderFn)
 
diff --git a/internal/clientcache/internal/daemon/options.go b/internal/clientcache/internal/daemon/options.go
index fca3930615..ba20725409 100644
--- a/internal/clientcache/internal/daemon/options.go
+++ b/internal/clientcache/internal/daemon/options.go
@@ -17,6 +17,7 @@ type options struct {
 	withRecheckSupportInterval             time.Duration
 	testWithIntervalRandomizationFactor    float64
 	testWithIntervalRandomizationFactorSet bool
+	WithReadyToServeNotificationCh         chan struct{}
 	withBoundaryTokenReaderFunc            cache.BoundaryTokenReaderFn
 
 	withUrl    string
@@ -99,3 +100,14 @@ func WithBoundaryTokenReaderFunc(_ context.Context, fn cache.BoundaryTokenReader
 		return nil
 	}
 }
+
+// WithReadyToServeNotificationCh provides an optional channel to notify when
+// the server is ready to serve; mainly used for test timing but exported for
+// availability. The channel will be closed just before the HTTP server is
+// started.
+func WithReadyToServeNotificationCh(_ context.Context, readyCh chan struct{}) Option {
+	return func(o *options) error {
+		o.WithReadyToServeNotificationCh = readyCh
+		return nil
+	}
+}
diff --git a/internal/clientcache/internal/daemon/options_test.go b/internal/clientcache/internal/daemon/options_test.go
index c87eef396f..f858238c2d 100644
--- a/internal/clientcache/internal/daemon/options_test.go
+++ b/internal/clientcache/internal/daemon/options_test.go
@@ -101,4 +101,14 @@ func Test_GetOpts(t *testing.T) {
 		testOpts := getDefaultOptions()
 		assert.Equal(t, opts, testOpts)
 	})
+	t.Run("WithReadyToServeNotificationCh", func(t *testing.T) {
+		ch := make(chan struct{})
+		opts, err := getOpts(WithReadyToServeNotificationCh(ctx, ch))
+		require.NoError(t, err)
+		assert.NotNil(t, opts.WithReadyToServeNotificationCh)
+		testOpts := getDefaultOptions()
+		assert.Nil(t, testOpts.WithReadyToServeNotificationCh)
+		testOpts.WithReadyToServeNotificationCh = ch
+		assert.Equal(t, opts, testOpts)
+	})
 }
diff --git a/internal/clientcache/internal/daemon/server.go b/internal/clientcache/internal/daemon/server.go
index 81aa8aaa88..79a80e67a8 100644
--- a/internal/clientcache/internal/daemon/server.go
+++ b/internal/clientcache/internal/daemon/server.go
@@ -338,6 +338,9 @@ func (s *CacheServer) Serve(ctx context.Context, cmd Commander, opt ...Option) e
 			return ctx
 		},
 	}
+	if opts.WithReadyToServeNotificationCh != nil {
+		close(opts.WithReadyToServeNotificationCh)
+	}
 	if err = s.httpSrv.Serve(l); err != nil && err != http.ErrServerClosed && !errors.Is(err, net.ErrClosed) {
 		event.WriteSysEvent(ctx, op, "error closing server", "err", err.Error())
 	}