From 3dd5760875047b19d719e5f3ab1559e2b38e40da Mon Sep 17 00:00:00 2001
From: "yash.bansal" <yash.bansal@redbus.com>
Date: Thu, 12 Sep 2024 12:58:18 +0530
Subject: [PATCH] Fix tests

---
 pubsub/tests/test_asserts.go | 2 +-
 pubsub/tests/test_pubsub.go  | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

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.