Merge a5593c1 into backport/jefferai-fix-tests/thankfully-electric-oc…

…elot
hashicorp · Sep 18, 2024 · a6f5b5a · a6f5b5a
2 parents 5d44819 + a5593c1
commit a6f5b5a
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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())
 	}