diff --git a/server/logging/nats_test.go b/server/logging/nats_test.go index 3a3c1e01ae..9c064fc7bb 100644 --- a/server/logging/nats_test.go +++ b/server/logging/nats_test.go @@ -83,6 +83,25 @@ func makeNatsLogs(t *testing.T) []json.RawMessage { return logs } +// natsWaitOrTimeout waits for the WaitGroup with a timeout to prevent +// tests from hanging indefinitely if messages are lost. +func natsWaitOrTimeout(t *testing.T, wg *sync.WaitGroup, timeout time.Duration) { + t.Helper() + + done := make(chan struct{}) + + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(timeout): + t.Fatal("timed out waiting for NATS messages") + } +} + func TestNatsLogRouter(t *testing.T) { // Define an abbreviated test query result. testLog := json.RawMessage(`{ @@ -163,6 +182,10 @@ func TestNatsLogWriter(t *testing.T) { // Ensure the subscription was created successfully. require.NoError(t, err) + // Flush to ensure the subscription is registered on the server + // before the writer publishes on its own connection. + require.NoError(t, nc.Flush()) + // Create the NATS log writer, specifying that the logs should be // published directly to the NATS subject, without using JetStream. writer, err := NewNatsLogWriter( @@ -185,7 +208,7 @@ func TestNatsLogWriter(t *testing.T) { require.NoError(t, writer.Write(t.Context(), expected)) // Wait for all logs to be received. - wg.Wait() + natsWaitOrTimeout(t, &wg, natsTestTimeout) // Ensure the received logs are equal to the expected logs. require.Equal(t, expected, received) @@ -382,6 +405,10 @@ func TestNatsLogWriter(t *testing.T) { // Ensure the subscription was created successfully. require.NoError(t, err) + // Flush to ensure the subscription is registered on the server + // before the writer publishes on its own connection. + require.NoError(t, nc.Flush()) + // Create the NATS log writer with gzip compression enabled. writer, err := NewNatsLogWriter( ns.ClientURL(), @@ -404,7 +431,7 @@ func TestNatsLogWriter(t *testing.T) { require.NoError(t, writer.Write(t.Context(), exp)) // Wait for all logs to be received. - wg.Wait() + natsWaitOrTimeout(t, &wg, natsTestTimeout) // Ensure the received logs are equal to the expected logs. require.Equal(t, exp, act) @@ -505,6 +532,10 @@ func TestNatsLogWriter(t *testing.T) { // Ensure the subscription was created successfully. require.NoError(t, err) + // Flush to ensure the subscription is registered on the server + // before the writer publishes on its own connection. + require.NoError(t, nc.Flush()) + // Create the NATS log writer with snappy compression enabled. writer, err := NewNatsLogWriter( ns.ClientURL(), @@ -527,7 +558,7 @@ func TestNatsLogWriter(t *testing.T) { require.NoError(t, writer.Write(t.Context(), exp)) // Wait for all logs to be received. - wg.Wait() + natsWaitOrTimeout(t, &wg, natsTestTimeout) // Ensure the received logs are equal to the expected logs. require.Equal(t, exp, act) @@ -567,6 +598,10 @@ func TestNatsLogWriter(t *testing.T) { // Ensure the subscription was created successfully. require.NoError(t, err) + // Flush to ensure the subscription is registered on the server + // before the writer publishes on its own connection. + require.NoError(t, nc.Flush()) + // Create the NATS log writer with zstd compression enabled. writer, err := NewNatsLogWriter( ns.ClientURL(), @@ -589,7 +624,7 @@ func TestNatsLogWriter(t *testing.T) { require.NoError(t, writer.Write(t.Context(), exp)) // Wait for all logs to be received. - wg.Wait() + natsWaitOrTimeout(t, &wg, natsTestTimeout) // Ensure the received logs are equal to the expected logs. require.Equal(t, exp, act)