From 3f2b671fb3c8da0c6c4e9c1cecf0b6935f4a5fec Mon Sep 17 00:00:00 2001
From: Nico <32375741+nulmete@users.noreply.github.com>
Date: Fri, 6 Feb 2026 10:28:40 -0300
Subject: [PATCH] Add natsWaitOrTimeout helper to fail early if any of the NATS
 logging tests fail (#39337)

Might resolve
[this](https://github.com/fleetdm/fleet/actions/runs/21648872745/job/62407941749?pr=39201#step:14:14722)

The writer could fire off messages before the subscriber was actually
registered on the NATS server (two separate connections, so no ordering
guarantee).
`nc.Flush()` forces a round-trip to make sure the subscription is in
place before we publish, and `natsWaitOrTimeout` is just a safety net so
we fail in 5s instead of hanging for 20min if something fails.
---
 server/logging/nats_test.go | 43 +++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

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)