cli/command: Add WithContextResolver to override context selection

Callers that choose the context from configuration had to do it before
Initialize. That meant loading the config file and context store
themselves and handling --config a second time.

WithContextResolver runs a callback inside Initialize after the config
file and context store are loaded, but before telemetry and endpoint
initialization. A non-empty result replaces the selected context, an
empty result keeps normal selection, and an error aborts
initialization.

The API client and endpoint are resolved lazily, so they use the
overridden context.

Signed-off-by: Paweł Gronowski <pawel.gronowski@docker.com>

cli/command/cli.go

@@ -68,6 +68,7 @@ type DockerCli struct {
 	serverInfo         ServerInfo
 	contextStore       store.Store
 	currentContext     string
+	contextResolver    func(*DockerCli) (string, error)
 	init               sync.Once
 	initErr            error
 	dockerEndpoint     docker.Endpoint
@@ -208,6 +209,8 @@ func (cli *DockerCli) HooksEnabled() bool {
 
 // Initialize the dockerCli runs initialization that must happen after command
 // line flags are parsed.
+//
+//nolint:gocyclo
 func (cli *DockerCli) Initialize(opts *cliflags.ClientOptions, ops ...CLIOption) error {
 	for _, o := range ops {
 		if err := o(cli); err != nil {
@@ -257,6 +260,16 @@ func (cli *DockerCli) Initialize(opts *cliflags.ClientOptions, ops ...CLIOption)
 		},
 	}
 
+	if cli.contextResolver != nil {
+		contextName, err := cli.contextResolver(cli)
+		if err != nil {
+			return err
+		}
+		if contextName != "" {
+			cli.currentContext = contextName
+		}
+	}
+
 	// TODO(krissetto): pass ctx to the funcs instead of using this
 	if cli.enableGlobalMeter {
 		cli.createGlobalMeterProvider(cli.baseCtx)

cli/command/cli_options.go

@@ -84,6 +84,19 @@ func WithDefaultContextStoreConfig() CLIOption {
 	}
 }
 
+// WithContextResolver overrides context selection during [DockerCli.Initialize].
+// The resolver runs after configuration and the context store are loaded, but
+// before telemetry and endpoint initialization.
+// An empty result preserves normal context selection; an error aborts
+// initialization.
+// The resolver must not initialize the API client.
+func WithContextResolver(resolve func(*DockerCli) (string, error)) CLIOption {
+	return func(cli *DockerCli) error {
+		cli.contextResolver = resolve
+		return nil
+	}
+}
+
 // WithAPIClient configures the cli to use the given API client.
 func WithAPIClient(c client.APIClient) CLIOption {
 	return func(cli *DockerCli) error {

cli/command/cli_test.go

@@ -325,6 +325,64 @@ func TestInitializeShouldAlwaysCreateTheContextStore(t *testing.T) {
 	assert.Check(t, cli.ContextStore() != nil)
 }
 
+func TestInitializeContextResolver(t *testing.T) {
+	resolverErr := errors.New("resolver failed")
+
+	for _, tc := range []struct {
+		name    string
+		context string
+		err     error
+	}{
+		{name: "resolved context", context: "resolved"},
+		{name: "normal selection"},
+		{name: "resolver failure", err: resolverErr},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			originalConfigDir := config.Dir()
+			t.Cleanup(func() { config.SetDir(originalConfigDir) })
+			config.SetDir(t.TempDir())
+
+			configDir := t.TempDir()
+			cfg := configfile.New(filepath.Join(configDir, config.ConfigFileName))
+			cfg.Features = map[string]string{"cloud": "foobar"}
+			assert.NilError(t, cfg.Save())
+
+			cli, err := NewDockerCli()
+			assert.NilError(t, err)
+
+			var loadedConfig *configfile.ConfigFile
+			calls := 0
+
+			err = cli.Initialize(&flags.ClientOptions{ConfigDir: configDir, Context: "original"},
+				WithContextResolver(func(cli *DockerCli) (string, error) {
+					calls++
+					assert.Equal(t, config.Dir(), configDir)
+					loadedConfig = cli.ConfigFile()
+					assert.Equal(t, loadedConfig.Features["cloud"], "foobar")
+					assert.Assert(t, cli.ContextStore() != nil)
+					assert.Assert(t, cli.client == nil)
+					return tc.context, tc.err
+				}),
+			)
+			assert.Equal(t, calls, 1)
+			assert.Assert(t, cli.ConfigFile() == loadedConfig)
+			assert.Assert(t, cli.client == nil)
+
+			if tc.err != nil {
+				assert.ErrorIs(t, err, tc.err)
+				return
+			}
+			assert.NilError(t, err)
+
+			wantContext := tc.context
+			if wantContext == "" {
+				wantContext = "original"
+			}
+			assert.Equal(t, cli.CurrentContext(), wantContext)
+		})
+	}
+}
+
 func TestHooksEnabled(t *testing.T) {
 	t.Run("disabled by default", func(t *testing.T) {
 		// Make sure we don't depend on any existing ~/.docker/config.json