Skip to content

Commit 7b0e2b8

Browse files
committed
testutils: allow overriding DefaultTestDRPCOption with env variable
This allows enabling or disabling DRPC for unit tests without changing code or recompiling. Note that the environment variable will take precedence over server args, which take precedence over the package-level option. Example usage: `dev test pkg/server:server_test -- --test_env=COCKROACH_TEST_DRPC=false` Epic: none Release note: none
1 parent 86d6f0c commit 7b0e2b8

File tree

1 file changed

+32
-1
lines changed

1 file changed

+32
-1
lines changed

pkg/testutils/serverutils/test_server_shim.go

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,13 @@ const (
177177

178178
testTenantModeEnabledShared = "shared"
179179
testTenantModeEnabledExternal = "external"
180+
181+
// COCKROACH_TEST_DRPC controls the DRPC enablement mode for test servers.
182+
//
183+
// - disabled: disables DRPC; all inter-node connectivity will use gRPC only
184+
//
185+
// - enabled: enables DRPC for inter-node connectivity
186+
testDRPCEnabledEnvVar = "COCKROACH_TEST_DRPC"
180187
)
181188

182189
func testTenantDecisionFromEnvironment(
@@ -537,17 +544,41 @@ func WaitForTenantCapabilities(
537544
}
538545
}
539546

547+
// parseDefaultTestDRPCOptionFromEnv parses the COCKROACH_TEST_DRPC environment
548+
// variable and returns the corresponding DefaultTestDRPCOption. If the
549+
// environment variable is not set it returns TestDRPCUnset. For invalid value,
550+
// it panic.
551+
func parseDefaultTestDRPCOptionFromEnv() base.DefaultTestDRPCOption {
552+
if str, present := envutil.EnvString(testDRPCEnabledEnvVar, 0); present {
553+
switch str {
554+
case "disabled", "false":
555+
return base.TestDRPCDisabled
556+
case "enabled", "true":
557+
return base.TestDRPCEnabled
558+
default:
559+
panic(fmt.Sprintf("invalid value for %s: %s", testDRPCEnabledEnvVar, str))
560+
}
561+
}
562+
return base.TestDRPCUnset
563+
}
564+
540565
// ShouldEnableDRPC determines the final DRPC option based on the input
541566
// option and any global overrides, resolving random choices to a concrete
542567
// enabled/disabled state.
543568
func ShouldEnableDRPC(
544569
ctx context.Context, t TestLogger, option base.DefaultTestDRPCOption,
545570
) base.DefaultTestDRPCOption {
546571
var logSuffix string
547-
if option == base.TestDRPCUnset && globalDefaultDRPCOptionOverride.isSet {
572+
573+
// Check environment variable first
574+
if envOption := parseDefaultTestDRPCOptionFromEnv(); envOption != base.TestDRPCUnset {
575+
option = envOption
576+
logSuffix = " (override by COCKROACH_TEST_DRPC environment variable)"
577+
} else if option == base.TestDRPCUnset && globalDefaultDRPCOptionOverride.isSet {
548578
option = globalDefaultDRPCOptionOverride.value
549579
logSuffix = " (override by TestingGlobalDRPCOption)"
550580
}
581+
551582
enableDRPC := false
552583
switch option {
553584
case base.TestDRPCEnabled:

0 commit comments

Comments
 (0)