Pass ctx explicitly

tests/integration/acp.go

@@ -112,7 +112,7 @@ func addPolicyACP(
 
 	nodeIDs, nodes := getNodesWithIDs(action.NodeID, s.nodes)
 	for index, node := range nodes {
-		ctx := getContextWithIdentity(s, action.Identity, nodeIDs[index])
+		ctx := getContextWithIdentity(s.ctx, s, action.Identity, nodeIDs[index])
 		policyResult, err := node.AddPolicy(ctx, action.Policy)
 
 		expectedErrorRaised := AssertError(s.t, s.testCase.Description, err, action.ExpectedError)
@@ -188,7 +188,7 @@ func addDocActorRelationshipACP(
 		collectionName, docID := getCollectionAndDocInfo(s, action.CollectionID, action.DocID, nodeID)
 
 		exists, err := node.AddDocActorRelationship(
-			getContextWithIdentity(s, action.RequestorIdentity, nodeID),
+			getContextWithIdentity(s.ctx, s, action.RequestorIdentity, nodeID),
 			collectionName,
 			docID,
 			action.Relation,
@@ -269,7 +269,7 @@ func deleteDocActorRelationshipACP(
 		collectionName, docID := getCollectionAndDocInfo(s, action.CollectionID, action.DocID, nodeID)
 
 		deleteDocActorRelationshipResult, err := node.DeleteDocActorRelationship(
-			getContextWithIdentity(s, action.RequestorIdentity, nodeID),
+			getContextWithIdentity(s.ctx, s, action.RequestorIdentity, nodeID),
 			collectionName,
 			docID,
 			action.Relation,

tests/integration/identity.go

@@ -15,10 +15,11 @@ import (
 	"math/rand"
 
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
-	acpIdentity "github.com/sourcenetwork/defradb/acp/identity"
-	identity "github.com/sourcenetwork/defradb/acp/identity"
 	"github.com/sourcenetwork/immutable"
 	"github.com/stretchr/testify/require"
+
+	acpIdentity "github.com/sourcenetwork/defradb/acp/identity"
+	identity "github.com/sourcenetwork/defradb/acp/identity"
 )
 
 // identRef is a type that refers to a specific identity of a certain type.
@@ -128,12 +129,12 @@ func generateIdentity(s *state) acpIdentity.Identity {
 // getContextWithIdentity returns a context with the identity for the given reference and node index.
 // If the identity does not exist, it will be generated.
 // The identity added to the context is prepared for a request, i.e. its [Identity.BearerToken] is set.
-func getContextWithIdentity(s *state, ref identRef, nodeIndex int) context.Context {
+func getContextWithIdentity(ctx context.Context, s *state, ref identRef, nodeIndex int) context.Context {
 	if !ref.hasValue {
-		return s.ctx
+		return ctx
 	}
 	ident := getIdentityForRequest(s, ref, nodeIndex)
-	return identity.WithContext(s.ctx, immutable.Some(ident))
+	return identity.WithContext(ctx, immutable.Some(ident))
 }
 
 func getIdentityDID(s *state, ident identRef) string {

tests/integration/state.go

@@ -129,7 +129,7 @@ type state struct {
 	txns []datastore.Txn
 
 	// identities contains all identities created in this test.
-	// The map key is the identity reference that uniquely identifies identities of different 
+	// The map key is the identity reference that uniquely identifies identities of different
 	// types. See [identRef].
 	// The map value is the identity holder that contains the identity itself and token
 	// generated for different target nodes. See [identityHolder].

tests/integration/utils.go

@@ -1313,7 +1313,7 @@ func createDocViaColSave(
 }
 
 func makeContextForDocCreate(s *state, ctx context.Context, nodeIndex int, action *CreateDoc) context.Context {
-	ctx = getContextWithIdentity(s, action.Identity, nodeIndex)
+	ctx = getContextWithIdentity(ctx, s, action.Identity, nodeIndex)
 	ctx = encryption.SetContextConfigFromParams(ctx, action.IsDocEncrypted, action.EncryptedFields)
 	return ctx
 }
@@ -1392,8 +1392,7 @@ func createDocViaGQL(
 	req := fmt.Sprintf(`mutation { %s(%s) { _docID } }`, key, params)
 
 	txn := getTransaction(s, node, immutable.None[int](), action.ExpectedError)
-	s.ctx = db.SetContextTxn(s.ctx, txn)
-	ctx := getContextWithIdentity(s, action.Identity, nodeIndex)
+	ctx := getContextWithIdentity(db.SetContextTxn(s.ctx, txn), s, action.Identity, nodeIndex)
 
 	result := node.ExecRequest(ctx, req)
 	if len(result.GQL.Errors) > 0 {
@@ -1447,7 +1446,7 @@ func deleteDoc(
 	for index, node := range nodes {
 		nodeID := nodeIDs[index]
 		collection := s.collections[nodeID][action.CollectionID]
-		ctx := getContextWithIdentity(s, action.Identity, nodeID)
+		ctx := getContextWithIdentity(s.ctx, s, action.Identity, nodeID)
 		err := withRetryOnNode(
 			node,
 			func() error {
@@ -1520,7 +1519,7 @@ func updateDocViaColSave(
 	nodeIndex int,
 	collection client.Collection,
 ) error {
-	ctx := getContextWithIdentity(s, action.Identity, nodeIndex)
+	ctx := getContextWithIdentity(s.ctx, s, action.Identity, nodeIndex)
 
 	doc, err := collection.Get(ctx, s.docIDs[action.CollectionID][action.DocID], true)
 	if err != nil {
@@ -1540,7 +1539,7 @@ func updateDocViaColUpdate(
 	nodeIndex int,
 	collection client.Collection,
 ) error {
-	ctx := getContextWithIdentity(s, action.Identity, nodeIndex)
+	ctx := getContextWithIdentity(s.ctx, s, action.Identity, nodeIndex)
 
 	doc, err := collection.Get(ctx, s.docIDs[action.CollectionID][action.DocID], true)
 	if err != nil {
@@ -1576,7 +1575,7 @@ func updateDocViaGQL(
 		input,
 	)
 
-	ctx := getContextWithIdentity(s, action.Identity, nodeIndex)
+	ctx := getContextWithIdentity(s.ctx, s, action.Identity, nodeIndex)
 
 	result := node.ExecRequest(ctx, request)
 	if len(result.GQL.Errors) > 0 {
@@ -1594,7 +1593,7 @@ func updateWithFilter(s *state, action UpdateWithFilter) {
 	for index, node := range nodes {
 		nodeID := nodeIDs[index]
 		collection := s.collections[nodeID][action.CollectionID]
-		ctx := getContextWithIdentity(s, action.Identity, nodeID)
+		ctx := getContextWithIdentity(s.ctx, s, action.Identity, nodeID)
 		err := withRetryOnNode(
 			node,
 			func() error {
@@ -1835,8 +1834,7 @@ func executeRequest(
 		nodeID := nodeIDs[index]
 		txn := getTransaction(s, node, action.TransactionID, action.ExpectedError)
 
-		s.ctx = db.SetContextTxn(s.ctx, txn)
-		ctx := getContextWithIdentity(s, action.Identity, nodeID)
+		ctx := getContextWithIdentity(db.SetContextTxn(s.ctx, txn), s, action.Identity, nodeID)
 
 		var options []client.RequestOption
 		if action.OperationName.HasValue() {