Fix tests

ThreeDotsLabs · Sep 12, 2024 · 9a6c148 · 9a6c148
1 parent 1447c26
commit 9a6c148
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/pubsub/tests/test_asserts.go b/pubsub/tests/test_asserts.go
@@ -99,7 +99,7 @@ func AssertAllMessagesHaveSameContext(t *testing.T, contextKeyString string, exp
 	assert.Len(t, received, len(expectedValues))
 	for _, msg := range received {
 		expectedValue := expectedValues[msg.UUID].Value(contextKey(contextKeyString)).(string)
-		actualValue := msg.Context().Value(contextKeyString)
+		actualValue := msg.Context().Value(contextKey(contextKeyString))
 		assert.Equal(t, expectedValue, actualValue)
 	}
 }
diff --git a/pubsub/tests/test_pubsub.go b/pubsub/tests/test_pubsub.go
@@ -1044,7 +1044,9 @@ ClosedLoop:
 
 	receivedMessages, _ := bulkRead(tCtx, msgs, messagesCount, defaultTimeout)
 	AssertAllMessagesReceived(t, publishedMessages, receivedMessages)
-	AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMessages)
+	if tCtx.Features.ContextPreserved {
+		AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMessages)
+	}
 }
 
 // TestReconnect tests if reconnecting to a Pub/Sub works correctly.