From 652236b36bbbd1996b014f7c4e80e33a74b9231b Mon Sep 17 00:00:00 2001
From: Christian Jackson-Gruber <christian.gruber@smartcontract.com>
Date: Mon, 25 Nov 2024 13:44:10 -0800
Subject: [PATCH] Add context.Context to deployment.Environment. Plumb this
 through in the test environment and devenv also. Mostly the ctx was already
 present for other reasons, but now we expose it through the environment for
 attaching to remote-calls that aren't already managed by a wrapper client.

---
 deployment/ccip/changeset/test_helpers.go    | 3 ++-
 deployment/environment.go                    | 3 +++
 deployment/environment/devenv/environment.go | 1 +
 deployment/environment/memory/environment.go | 9 +++++++--
 deployment/keystone/deploy_test.go           | 6 ++++--
 5 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/deployment/ccip/changeset/test_helpers.go b/deployment/ccip/changeset/test_helpers.go
index 5e8705d4757..b46a42d7556 100644
--- a/deployment/ccip/changeset/test_helpers.go
+++ b/deployment/ccip/changeset/test_helpers.go
@@ -16,6 +16,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/pkg/errors"
 	"github.com/smartcontractkit/ccip-owner-contracts/pkg/gethwrappers"
+
 	"github.com/smartcontractkit/chainlink-ccip/pluginconfig"
 
 	commonconfig "github.com/smartcontractkit/chainlink-common/pkg/config"
@@ -200,7 +201,7 @@ func NewMemoryEnvironment(
 			require.NoError(t, node.App.Stop())
 		})
 	}
-	e := memory.NewMemoryEnvironmentFromChainsNodes(t, lggr, chains, nodes)
+	e := memory.NewMemoryEnvironmentFromChainsNodes(lggr, ctx, chains, nodes)
 	envNodes, err := deployment.NodeInfo(e.NodeIDs, e.Offchain)
 	require.NoError(t, err)
 	e.ExistingAddresses = ab
diff --git a/deployment/environment.go b/deployment/environment.go
index c55f4d3efe6..2e33df2ffcc 100644
--- a/deployment/environment.go
+++ b/deployment/environment.go
@@ -76,6 +76,7 @@ type Environment struct {
 	Chains            map[uint64]Chain
 	NodeIDs           []string
 	Offchain          OffchainClient
+	Context           context.Context
 }
 
 func NewEnvironment(
@@ -85,6 +86,7 @@ func NewEnvironment(
 	chains map[uint64]Chain,
 	nodeIDs []string,
 	offchain OffchainClient,
+	context context.Context,
 ) *Environment {
 	return &Environment{
 		Name:              name,
@@ -93,6 +95,7 @@ func NewEnvironment(
 		Chains:            chains,
 		NodeIDs:           nodeIDs,
 		Offchain:          offchain,
+		Context:           context,
 	}
 }
 
diff --git a/deployment/environment/devenv/environment.go b/deployment/environment/devenv/environment.go
index 94319d2247e..51a5a6bc231 100644
--- a/deployment/environment/devenv/environment.go
+++ b/deployment/environment/devenv/environment.go
@@ -53,5 +53,6 @@ func NewEnvironment(ctx context.Context, lggr logger.Logger, config EnvironmentC
 		chains,
 		nodeIDs,
 		offChain,
+		ctx,
 	), jd.don, nil
 }
diff --git a/deployment/environment/memory/environment.go b/deployment/environment/memory/environment.go
index a1478a3bf52..b5fa31d83ab 100644
--- a/deployment/environment/memory/environment.go
+++ b/deployment/environment/memory/environment.go
@@ -13,6 +13,7 @@ import (
 
 	chainsel "github.com/smartcontractkit/chain-selectors"
 
+	"github.com/smartcontractkit/chainlink-testing-framework/lib/utils/testcontext"
 	"github.com/smartcontractkit/chainlink/deployment"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
@@ -110,10 +111,12 @@ func NewNodes(t *testing.T, logLevel zapcore.Level, chains map[uint64]deployment
 	return nodesByPeerID
 }
 
-func NewMemoryEnvironmentFromChainsNodes(t *testing.T,
+func NewMemoryEnvironmentFromChainsNodes(
 	lggr logger.Logger,
+	ctx context.Context,
 	chains map[uint64]deployment.Chain,
-	nodes map[string]Node) deployment.Environment {
+	nodes map[string]Node,
+) deployment.Environment {
 	var nodeIDs []string
 	for id := range nodes {
 		nodeIDs = append(nodeIDs, id)
@@ -125,6 +128,7 @@ func NewMemoryEnvironmentFromChainsNodes(t *testing.T,
 		chains,
 		nodeIDs, // Note these have the p2p_ prefix.
 		NewMemoryJobClient(nodes),
+		ctx,
 	)
 }
 
@@ -143,5 +147,6 @@ func NewMemoryEnvironment(t *testing.T, lggr logger.Logger, logLevel zapcore.Lev
 		chains,
 		nodeIDs,
 		NewMemoryJobClient(nodes),
+		testcontext.Get(t),
 	)
 }
diff --git a/deployment/keystone/deploy_test.go b/deployment/keystone/deploy_test.go
index e446405944c..7a70fad5d86 100644
--- a/deployment/keystone/deploy_test.go
+++ b/deployment/keystone/deploy_test.go
@@ -9,8 +9,10 @@ import (
 
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	chainsel "github.com/smartcontractkit/chain-selectors"
+
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
 	"github.com/smartcontractkit/chainlink-common/pkg/utils/tests"
+	"github.com/smartcontractkit/chainlink-testing-framework/lib/utils/testcontext"
 	"github.com/smartcontractkit/chainlink/deployment"
 	"github.com/smartcontractkit/chainlink/deployment/environment/clo"
 	"github.com/smartcontractkit/chainlink/deployment/environment/clo/models"
@@ -26,6 +28,7 @@ import (
 
 func TestDeploy(t *testing.T) {
 	lggr := logger.Test(t)
+	ctx := testcontext.Get(t)
 
 	// sepolia; all nodes are on the this chain
 	sepoliaChainId := uint64(11155111)
@@ -100,7 +103,7 @@ func TestDeploy(t *testing.T) {
 	maps.Copy(allNodes, wfNodes)
 	maps.Copy(allNodes, cwNodes)
 	maps.Copy(allNodes, assetNodes)
-	env := memory.NewMemoryEnvironmentFromChainsNodes(t, lggr, allChains, allNodes)
+	env := memory.NewMemoryEnvironmentFromChainsNodes(lggr, ctx, allChains, allNodes)
 
 	var ocr3Config = keystone.OracleConfigWithSecrets{
 		OracleConfig: keystone.OracleConfig{
@@ -109,7 +112,6 @@ func TestDeploy(t *testing.T) {
 		OCRSecrets: deployment.XXXGenerateTestOCRSecrets(),
 	}
 
-	ctx := tests.Context(t)
 	// explicitly deploy the contracts
 	cs, err := keystone.DeployContracts(lggr, &env, sepoliaChainSel)
 	require.NoError(t, err)