Refactor SimualtedEnvironment to builder pattern

cmd/workflow/deploy/deploy_test.go

@@ -431,10 +431,10 @@ func TestConfigFlagsMutuallyExclusive(t *testing.T) {
 
 func TestValidateInputs_PrivateRegistry(t *testing.T) {
 	t.Run("accepts URL wasm and config with off-chain resolved registry and no on-chain contract inputs", func(t *testing.T) {
-		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
+		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t).WithPrivateRegistry("42", "test_label")
 		defer simulatedEnvironment.Close()
 
-		ctx, buf := simulatedEnvironment.NewOffChainRuntimeContextWithBufferedOutput("42", "test_label")
+		ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
 		h := newHandler(ctx, buf)
 		ctx.Settings = createTestSettings(
 			chainsim.TestAddress,
@@ -464,10 +464,10 @@ func TestValidateInputs_PrivateRegistry(t *testing.T) {
 	})
 
 	t.Run("fails when required don family is missing for private target", func(t *testing.T) {
-		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
+		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t).WithPrivateRegistry("42", "")
 		defer simulatedEnvironment.Close()
 
-		ctx, buf := simulatedEnvironment.NewOffChainRuntimeContextWithBufferedOutput("42", "")
+		ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
 		h := newHandler(ctx, buf)
 		ctx.Settings = createTestSettings(
 			chainsim.TestAddress,
@@ -740,10 +740,10 @@ func containsQuery(query, operation string) bool {
 
 func newPrivateRegistryExecuteHandler(t *testing.T, wasmURL, gqlURL string) *handler {
 	t.Helper()
-	simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
+	simulatedEnvironment := chainsim.NewSimulatedEnvironment(t).WithPrivateRegistry("", "test-don")
 	t.Cleanup(simulatedEnvironment.Close)
 
-	ctx, buf := simulatedEnvironment.NewOffChainRuntimeContextWithBufferedOutput("", "test-don")
+	ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
 	h := newHandler(ctx, buf)
 	ctx.Settings = createTestSettings(
 		chainsim.TestAddress,

cmd/workflow/deploy/private_registry_test.go

@@ -321,7 +321,7 @@ func TestResolveWorkflowOwner(t *testing.T) {
 		require.NoError(t, err)
 		expectedOwner := "0x" + hex.EncodeToString(expectedBytes)
 
-		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
+		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t).WithPrivateRegistry("42", "test-don")
 		defer simulatedEnvironment.Close()
 
 		ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
@@ -344,7 +344,7 @@ func TestResolveWorkflowOwner(t *testing.T) {
 	t.Run("private target adds 0x prefix when missing", func(t *testing.T) {
 		t.Parallel()
 
-		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
+		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t).WithPrivateRegistry("42", "test-don")
 		defer simulatedEnvironment.Close()
 
 		ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
@@ -366,7 +366,7 @@ func TestResolveWorkflowOwner(t *testing.T) {
 
 	t.Run("private target errors when derived workflow owner is empty", func(t *testing.T) {
 		t.Parallel()
-		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
+		simulatedEnvironment := chainsim.NewSimulatedEnvironment(t).WithPrivateRegistry("42", "test-don")
 		defer simulatedEnvironment.Close()
 
 		ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()

internal/credentials/credentials_test.go

@@ -1,14 +1,13 @@
 package credentials
 
 import (
-	"encoding/base64"
-	"encoding/json"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 
 	"github.com/smartcontractkit/cre-cli/internal/testutil"
+	"github.com/smartcontractkit/cre-cli/internal/testutil/testjwt"
 )
 
 func TestNew_Default(t *testing.T) {
@@ -86,31 +85,13 @@ TokenType:  "file-type"
 	}
 }
 
-// Helper function to create a JWT token with custom claims
-func createTestJWT(t *testing.T, claims map[string]interface{}) string {
-	t.Helper()
-
-	// JWT header (doesn't matter for our tests)
-	header := map[string]string{"alg": "HS256", "typ": "JWT"}
-	headerJSON, _ := json.Marshal(header)
-	headerEncoded := base64.RawURLEncoding.EncodeToString(headerJSON)
-
-	// JWT payload with claims
-	claimsJSON, err := json.Marshal(claims)
-	if err != nil {
-		t.Fatalf("failed to marshal claims: %v", err)
-	}
-	claimsEncoded := base64.RawURLEncoding.EncodeToString(claimsJSON)
-
-	// JWT signature (doesn't need to be valid for our tests)
-	signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
-
-	return headerEncoded + "." + claimsEncoded + "." + signature
+func createTestJWT(claims map[string]interface{}) string {
+	return testjwt.CreateTestJWTWithClaims(claims)
 }
 
 func TestGetOrgID_BearerWithOrgID(t *testing.T) {
 	logger := testutil.NewTestLogger()
-	token := createTestJWT(t, map[string]interface{}{
+	token := createTestJWT(map[string]interface{}{
 		"sub":    "user123",
 		"org_id": "org_abc123",
 	})
@@ -132,7 +113,7 @@ func TestGetOrgID_BearerWithOrgID(t *testing.T) {
 
 func TestGetOrgID_MissingClaim(t *testing.T) {
 	logger := testutil.NewTestLogger()
-	token := createTestJWT(t, map[string]interface{}{
+	token := createTestJWT(map[string]interface{}{
 		"sub": "user123",
 	})
 
@@ -153,7 +134,7 @@ func TestGetOrgID_MissingClaim(t *testing.T) {
 
 func TestGetOrgID_EmptyClaim(t *testing.T) {
 	logger := testutil.NewTestLogger()
-	token := createTestJWT(t, map[string]interface{}{
+	token := createTestJWT(map[string]interface{}{
 		"sub":    "user123",
 		"org_id": "",
 	})
@@ -260,7 +241,7 @@ func TestCheckIsUngatedOrganization_JWTWithFullAccess(t *testing.T) {
 				tc.namespace + "organization_roles":  "ROOT",
 			}
 
-			token := createTestJWT(t, claims)
+			token := createTestJWT(claims)
 
 			creds := &Credentials{
 				AuthType: AuthTypeBearer,
@@ -289,7 +270,7 @@ func TestCheckIsUngatedOrganization_JWTWithMissingClaim(t *testing.T) {
 		// organization_status claim is missing
 	}
 
-	token := createTestJWT(t, claims)
+	token := createTestJWT(claims)
 
 	creds := &Credentials{
 		AuthType: AuthTypeBearer,
@@ -317,7 +298,7 @@ func TestCheckIsUngatedOrganization_JWTWithEmptyStatus(t *testing.T) {
 		"https://api.cre.chain.link/organization_status": "",
 	}
 
-	token := createTestJWT(t, claims)
+	token := createTestJWT(claims)
 
 	creds := &Credentials{
 		AuthType: AuthTypeBearer,
@@ -345,7 +326,7 @@ func TestCheckIsUngatedOrganization_JWTWithGatedStatus(t *testing.T) {
 		"https://api.cre.chain.link/organization_status": "GATED",
 	}
 
-	token := createTestJWT(t, claims)
+	token := createTestJWT(claims)
 
 	creds := &Credentials{
 		AuthType: AuthTypeBearer,
@@ -373,7 +354,7 @@ func TestCheckIsUngatedOrganization_JWTWithRestrictedStatus(t *testing.T) {
 		"https://api.cre.chain.link/organization_status": "RESTRICTED",
 	}
 
-	token := createTestJWT(t, claims)
+	token := createTestJWT(claims)
 
 	creds := &Credentials{
 		AuthType: AuthTypeBearer,

internal/testutil/chainsim/simulated_environment.go

@@ -15,13 +15,18 @@ import (
 	settingspkg "github.com/smartcontractkit/cre-cli/internal/settings"
 	"github.com/smartcontractkit/cre-cli/internal/tenantctx"
 	"github.com/smartcontractkit/cre-cli/internal/testutil"
+	"github.com/smartcontractkit/cre-cli/internal/testutil/testjwt"
 	"github.com/smartcontractkit/cre-cli/internal/testutil/testsettings"
 )
 
 type SimulatedEnvironment struct {
 	Chain     *SimulatedChain
 	EthClient *seth.Client
 	Contracts *SimulatedContracts
+
+	tenantID  string
+	donFamily string
+	jwtToken  string
 }
 
 type SimulatedContracts struct {
@@ -47,6 +52,17 @@ func NewSimulatedEnvironment(t *testing.T) *SimulatedEnvironment {
 	return &simulatedEnvironment
 }
 
+func (se *SimulatedEnvironment) WithPrivateRegistry(tenantID, donFamily string) *SimulatedEnvironment {
+	se.tenantID = tenantID
+	se.donFamily = donFamily
+	return se
+}
+
+func (se *SimulatedEnvironment) WithJWT(orgID string) *SimulatedEnvironment {
+	se.jwtToken = testjwt.CreateTestJWT(orgID)
+	return se
+}
+
 func (se *SimulatedEnvironment) NewRuntimeContext() *runtime.Context {
 	logger := testutil.NewTestLogger()
 	return se.createContextWithLogger(logger)
@@ -57,20 +73,6 @@ func (se *SimulatedEnvironment) NewRuntimeContextWithBufferedOutput() (*runtime.
 	return se.createContextWithLogger(logger), buf
 }
 
-func (se *SimulatedEnvironment) NewOffChainRuntimeContext(tenantID, donFamily string) *runtime.Context {
-	ctx := se.NewRuntimeContext()
-	ctx.TenantContext = &tenantctx.EnvironmentContext{TenantID: tenantID}
-	ctx.ResolvedRegistry = settingspkg.NewOffChainRegistry("private", donFamily)
-	return ctx
-}
-
-func (se *SimulatedEnvironment) NewOffChainRuntimeContextWithBufferedOutput(tenantID, donFamily string) (*runtime.Context, *bytes.Buffer) {
-	ctx, buf := se.NewRuntimeContextWithBufferedOutput()
-	ctx.TenantContext = &tenantctx.EnvironmentContext{TenantID: tenantID}
-	ctx.ResolvedRegistry = settingspkg.NewOffChainRegistry("private", donFamily)
-	return ctx, buf
-}
-
 func (se *SimulatedEnvironment) Close() {
 	se.Chain.Close()
 }
@@ -96,7 +98,9 @@ func (se *SimulatedEnvironment) createContextWithLogger(logger *zerolog.Logger)
 	}
 
 	var resolved settingspkg.ResolvedRegistry
-	if environmentSet != nil {
+	if se.tenantID != "" && se.donFamily != "" {
+		resolved = settingspkg.NewOffChainRegistry("private", se.donFamily)
+	} else if environmentSet != nil {
 		resolved = settingspkg.NewOnChainRegistry(
 			"",
 			se.Contracts.WorkflowRegistry.Contract.Hex(),
@@ -116,9 +120,19 @@ func (se *SimulatedEnvironment) createContextWithLogger(logger *zerolog.Logger)
 		ResolvedRegistry: resolved,
 	}
 
+	if se.tenantID != "" {
+		ctx.TenantContext = &tenantctx.EnvironmentContext{TenantID: se.tenantID}
+	}
+
 	// Mark credentials as validated for tests to bypass validation
 	if creds != nil {
 		creds.IsValidated = true
+		if se.jwtToken != "" {
+			if creds.Tokens == nil {
+				creds.Tokens = &credentials.CreLoginTokenSet{}
+			}
+			creds.Tokens.AccessToken = se.jwtToken
+		}
 	}
 
 	return ctx

internal/testutil/testjwt/jwt.go

@@ -0,0 +1,35 @@
+package testjwt
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"time"
+)
+
+// CreateTestJWTWithClaims creates a JWT token with custom claims for testing.
+// The signature is a dummy value.
+func CreateTestJWTWithClaims(claims map[string]interface{}) string {
+	// JWT header (doesn't matter for our tests)
+	header := map[string]string{"alg": "HS256", "typ": "JWT"}
+	headerJSON, _ := json.Marshal(header)
+	headerEncoded := base64.RawURLEncoding.EncodeToString(headerJSON)
+
+	// JWT payload with claims
+	claimsJSON, _ := json.Marshal(claims)
+	claimsEncoded := base64.RawURLEncoding.EncodeToString(claimsJSON)
+
+	// JWT signature (doesn't need to be valid for our tests)
+	signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
+
+	return headerEncoded + "." + claimsEncoded + "." + signature
+}
+
+// CreateTestJWT creates a JWT token with default claims for a given organization ID.
+func CreateTestJWT(orgID string) string {
+	return CreateTestJWTWithClaims(map[string]interface{}{
+		"sub":                 "test-user",
+		"org_id":              orgID,
+		"organization_status": "FULL_ACCESS",
+		"exp":                 time.Now().Add(2 * time.Hour).Unix(),
+	})
+}

test/multi_command_flows/workflow_private_registry.go

@@ -2,7 +2,6 @@ package multi_command_flows
 
 import (
 	"bytes"
-	"encoding/base64"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -12,12 +11,12 @@ import (
 	"strings"
 	"sync/atomic"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/require"
 
 	"github.com/smartcontractkit/cre-cli/internal/environments"
 	"github.com/smartcontractkit/cre-cli/internal/settings"
+	"github.com/smartcontractkit/cre-cli/internal/testutil/testjwt"
 )
 
 const privateRegistryOwnerAddress = "0x0bb6e890f43f93f4f4f5eb64fdf17f81f15bf12a"
@@ -42,18 +41,7 @@ func createTestBearerCredentialsHome(t *testing.T) string {
 }
 
 func createTestJWT(orgID string) string {
-	header, _ := json.Marshal(map[string]any{"alg": "none", "typ": "JWT"})
-	payload, _ := json.Marshal(map[string]any{
-		"sub":                 "test-user",
-		"org_id":              orgID,
-		"organization_status": "FULL_ACCESS",
-		"exp":                 time.Now().Add(2 * time.Hour).Unix(),
-	})
-
-	headerEnc := base64.RawURLEncoding.EncodeToString(header)
-	payloadEnc := base64.RawURLEncoding.EncodeToString(payload)
-
-	return headerEnc + "." + payloadEnc + ".signature"
+	return testjwt.CreateTestJWT(orgID)
 }
 
 // workflowDeployPrivateRegistry deploys a workflow to the private registry via CLI