From dd1d13d9c69ef8fcb2d4e6fa5bc1abfbda2ac275 Mon Sep 17 00:00:00 2001 From: Nukitt Date: Mon, 24 Nov 2025 21:22:54 +0530 Subject: [PATCH] testserver: consolidate testserver config to InitTestServerFactory Currently, we use global overrides to configure testserver's tenant and DRPC options. These were also using defer statements, but it was concluded that there is not much value in actually performing the cleanup function to reset the global state option in a defer, since the test process exits immediately after finishing the tests. This change introduces the framework changes required to consolidate the configuration into InitTestServerFactory. For now, we also have the global overrides present to maintain backward compatibility until we can have all the usages of the `TestingSetDefaultTenantSelectionOverride` and `TestingGlobalDRPCOption` removed. We also change the priority order of changing these configs, highest priority is given to the config set in the test itself, then the global override, then factory default and then the enviorment variable change through the cli. After this changes, we can simply add the configuration we want for the package into `InitTestServerFactory`. Example usage of it would look like: ```go serverutils.InitTestServerFactory(server.TestServerFactory, serverutils.WithDRPCOption(base.TestDRPCEnabled), serverutils.WithTenantOption(base.TestTenantProbabilisticOnly, ) ``` Epic: None Fixes: #158004 Release note: None --- pkg/testutils/serverutils/test_server_shim.go | 117 +++++++++++++----- 1 file changed, 85 insertions(+), 32 deletions(-) diff --git a/pkg/testutils/serverutils/test_server_shim.go b/pkg/testutils/serverutils/test_server_shim.go index 9444d19c7fb3..9629379348cc 100644 --- a/pkg/testutils/serverutils/test_server_shim.go +++ b/pkg/testutils/serverutils/test_server_shim.go @@ -88,7 +88,9 @@ func ShouldStartDefaultTestTenant( // Explicit case for disabling the default test tenant. if baseArg.TestTenantAlwaysDisabled() { if issueNum, label := baseArg.IssueRef(); issueNum != 0 { - t.Logf("cluster virtualization disabled due to issue: #%d (expected label: %s)", issueNum, label) + if t != nil { + t.Logf("cluster virtualization disabled due to issue: #%d (expected label: %s)", issueNum, label) + } } return baseArg } @@ -103,9 +105,45 @@ func ShouldStartDefaultTestTenant( // If the test tenant is explicitly enabled and a process mode selected, then // we are done. if !baseArg.TestTenantNoDecisionMade() { + if t != nil { + t.Log("using test tenant configuration from test explicit setting") + } return baseArg } + if globalDefaultSelectionOverride.isSet { + override := globalDefaultSelectionOverride.value + if override.TestTenantNoDecisionMade() { + panic("programming error: global override does not contain a final decision") + } + if override.TestTenantAlwaysDisabled() { + if issueNum, label := override.IssueRef(); issueNum != 0 && t != nil { + t.Logf("cluster virtualization disabled in global scope due to issue: #%d (expected label: %s)", issueNum, label) + } + } else { + if t != nil { + t.Log(defaultTestTenantMessage(override.SharedProcessMode()) + "\n(via override from TestingSetDefaultTenantSelectionOverride)") + } + } + return override + } + + if factoryDefaultTenant != nil { + defaultArg := *factoryDefaultTenant + // If factory default made a decision, return it. + if !defaultArg.TestTenantNoDecisionMade() { + if t != nil { + t.Log("using test tenant configuration from testserver factory defaults") + } + if defaultArg.TestTenantAlwaysDisabled() { + if issueNum, label := defaultArg.IssueRef(); issueNum != 0 && t != nil { + t.Logf("cluster virtualization disabled due to issue: #%d (expected label: %s)", issueNum, label) + } + } + return defaultArg + } + } + // Determine if the default test tenant should be run as a shared process. var shared bool switch { @@ -127,24 +165,11 @@ func ShouldStartDefaultTestTenant( if decision, override := testTenantDecisionFromEnvironment(baseArg, shared); override { if decision.TestTenantAlwaysEnabled() { - t.Log(defaultTestTenantMessage(decision.SharedProcessMode()) + "\n(override via COCKROACH_TEST_TENANT)") - } - return decision - } - - if globalDefaultSelectionOverride.isSet { - override := globalDefaultSelectionOverride.value - if override.TestTenantNoDecisionMade() { - panic("programming error: global override does not contain a final decision") - } - if override.TestTenantAlwaysDisabled() { - if issueNum, label := override.IssueRef(); issueNum != 0 { - t.Logf("cluster virtualization disabled in global scope due to issue: #%d (expected label: %s)", issueNum, label) + if t != nil { + t.Log(defaultTestTenantMessage(decision.SharedProcessMode()) + "\n(override via COCKROACH_TEST_TENANT)") } - } else { - t.Log(defaultTestTenantMessage(override.SharedProcessMode()) + "\n(override via TestingSetDefaultTenantSelectionOverride)") } - return override + return decision } // Note: we ask the metamorphic framework for a "disable" value, instead @@ -217,9 +242,6 @@ var globalDefaultDRPCOptionOverride struct { } // TestingGlobalDRPCOption sets the package-level DefaultTestDRPCOption. -// -// Note: This override will be superseded by any more specific options provided -// when starting the server or cluster. func TestingGlobalDRPCOption(v base.DefaultTestDRPCOption) func() { globalDefaultDRPCOptionOverride.isSet = true globalDefaultDRPCOptionOverride.value = v @@ -245,12 +267,32 @@ func TestingSetDefaultTenantSelectionOverride(v base.DefaultTestTenantOptions) f } var srvFactoryImpl TestServerFactory +var factoryDefaultDRPC *base.DefaultTestDRPCOption +var factoryDefaultTenant *base.DefaultTestTenantOptions + +type TestServerFactoryOption func() + +func WithTenantOption(opt base.DefaultTestTenantOptions) TestServerFactoryOption { + return func() { + factoryDefaultTenant = &opt + } +} +func WithDRPCOption(opt base.DefaultTestDRPCOption) TestServerFactoryOption { + return func() { + factoryDefaultDRPC = &opt + } +} // InitTestServerFactory should be called once to provide the implementation // of the service. It will be called from a xx_test package that can import the -// server package. -func InitTestServerFactory(impl TestServerFactory) { +// server package. Optional parameters can be passed to set default options: +// - WithDRPCOption: Sets the default DRPC mode for test servers +// - WithTenantOption: Sets the default tenant configuration for test servers +func InitTestServerFactory(impl TestServerFactory, opts ...TestServerFactoryOption) { srvFactoryImpl = impl + for _, opt := range opts { + opt() + } } // TestLogger is the minimal interface of testing.T that is used by @@ -281,6 +323,8 @@ func StartServerOnlyE(t TestLogger, params base.TestServerArgs) (TestServerInter allowAdditionalTenants := params.DefaultTestTenant.AllowAdditionalTenants() // Update the flags with the actual decisions for test configuration. + // Priority of these Should* functions: + // Test explicit value > global override > factory defaults > env vars > metamorphic. params.DefaultTestTenant = ShouldStartDefaultTestTenant(t, params.DefaultTestTenant) params.DefaultDRPCOption = ShouldEnableDRPC(ctx, t, params.DefaultDRPCOption) @@ -579,20 +623,27 @@ func parseDefaultTestDRPCOptionFromEnv() base.DefaultTestDRPCOption { } // ShouldEnableDRPC determines the final DRPC option based on the input -// option and any global overrides, resolving random choices to a concrete -// enabled/disabled state. +// option, resolving random choices to a concrete enabled/disabled state. func ShouldEnableDRPC( ctx context.Context, t TestLogger, option base.DefaultTestDRPCOption, ) base.DefaultTestDRPCOption { var logSuffix string - // Check environment variable first - if envOption := parseDefaultTestDRPCOptionFromEnv(); envOption != base.TestDRPCUnset { - option = envOption - logSuffix = " (override by COCKROACH_TEST_DRPC environment variable)" - } else if option == base.TestDRPCUnset && globalDefaultDRPCOptionOverride.isSet { - option = globalDefaultDRPCOptionOverride.value - logSuffix = " (override by TestingGlobalDRPCOption)" + if option == base.TestDRPCUnset { + if globalDefaultDRPCOptionOverride.isSet { + option = globalDefaultDRPCOptionOverride.value + logSuffix = " (via override by TestingGlobalDRPCOption)" + } else if factoryDefaultDRPC != nil { + option = *factoryDefaultDRPC + logSuffix = " (via testserver factory defaults)" + } else if envOption := parseDefaultTestDRPCOptionFromEnv(); envOption != base.TestDRPCUnset { + option = envOption + logSuffix = " (via COCKROACH_TEST_DRPC environment variable)" + } else { + return base.TestDRPCUnset + } + } else { + logSuffix = " (via test explicit setting)" } enableDRPC := false @@ -607,7 +658,9 @@ func ShouldEnableDRPC( } if enableDRPC { - t.Log("DRPC is enabled" + logSuffix) + if t != nil { + t.Log("DRPC is enabled" + logSuffix) + } return base.TestDRPCEnabled }