Fix tests

yash.bansal · yash.bansal · commit 9a6c148d6e8e · 2024-09-12T12:58:18.000+05:30
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.

Original file line number	Diff line number	Diff line change
`@@ -99,7 +99,7 @@ func AssertAllMessagesHaveSameContext(t *testing.T, contextKeyString string, exp`
`99`	`99`	`assert.Len(t, received, len(expectedValues))`
`100`	`100`	`for _, msg := range received {`
`101`	`101`	`expectedValue := expectedValues[msg.UUID].Value(contextKey(contextKeyString)).(string)`
`102`		`- actualValue := msg.Context().Value(contextKeyString)`
	`102`	`+ actualValue := msg.Context().Value(contextKey(contextKeyString))`
`103`	`103`	`assert.Equal(t, expectedValue, actualValue)`
`104`	`104`	`}`
`105`	`105`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1044,7 +1044,9 @@ ClosedLoop:`
`1044`	`1044`
`1045`	`1045`	`receivedMessages, _ := bulkRead(tCtx, msgs, messagesCount, defaultTimeout)`
`1046`	`1046`	`AssertAllMessagesReceived(t, publishedMessages, receivedMessages)`
`1047`		`- AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMessages)`
	`1047`	`+ if tCtx.Features.ContextPreserved {`
	`1048`	`+ AssertAllMessagesHaveSameContext(t, contextKeyString, expectedContexts, receivedMessages)`
	`1049`	`+ }`
`1048`	`1050`	`}`
`1049`	`1051`
`1050`	`1052`	`// TestReconnect tests if reconnecting to a Pub/Sub works correctly.`