diff --git a/pubsub/tests/test_asserts.go b/pubsub/tests/test_asserts.go
index 5942b2838..7f2798b92 100644
--- 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
index 1a5566c01..ac6efbdba 100644
--- 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.