From fdd2758caa32d8447bb33c5ee62d38e17e602f02 Mon Sep 17 00:00:00 2001
From: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
Date: Sat, 15 Nov 2025 17:09:18 -0500
Subject: [PATCH 1/4] standardize workstation tests

Signed-off-by: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
---
 pkg/workstation/workstation_test.go | 91 +++++++++++++----------------
 1 file changed, 42 insertions(+), 49 deletions(-)

diff --git a/pkg/workstation/workstation_test.go b/pkg/workstation/workstation_test.go
index 30670e9f5..74c95b3d9 100644
--- a/pkg/workstation/workstation_test.go
+++ b/pkg/workstation/workstation_test.go
@@ -21,7 +21,7 @@ import (
 // Test Setup
 // =============================================================================
 
-type Mocks struct {
+type WorkstationTestMocks struct {
 	Runtime          *ctxpkg.Runtime
 	ConfigHandler    config.ConfigHandler
 	Shell            *shell.MockShell
@@ -32,11 +32,6 @@ type Mocks struct {
 	SSHClient        *ssh.MockClient
 }
 
-type SetupOptions struct {
-	ConfigHandler config.ConfigHandler
-	ConfigStr     string
-}
-
 func convertToServiceSlice(mockServices []*services.MockService) []services.Service {
 	serviceSlice := make([]services.Service, len(mockServices))
 	for i, mockService := range mockServices {
@@ -45,7 +40,7 @@ func convertToServiceSlice(mockServices []*services.MockService) []services.Serv
 	return serviceSlice
 }
 
-func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupWorkstationMocks(t *testing.T, opts ...func(*WorkstationTestMocks)) *WorkstationTestMocks {
 	t.Helper()
 
 	// Create mock config handler
@@ -176,15 +171,6 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 	mockContainerRuntime.UpFunc = func(verbose ...bool) error { return nil }
 	mockContainerRuntime.DownFunc = func() error { return nil }
 
-	// Apply custom options
-	if len(opts) > 0 && opts[0] != nil {
-		if opts[0].ConfigHandler != nil {
-			if mockConfig, ok := opts[0].ConfigHandler.(*config.MockConfigHandler); ok {
-				mockConfigHandler = mockConfig
-			}
-		}
-	}
-
 	rt := &ctxpkg.Runtime{
 		ContextName:   "test-context",
 		ProjectRoot:   "/test/project",
@@ -194,7 +180,7 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 		Shell:         mockShell,
 	}
 
-	return &Mocks{
+	mocks := &WorkstationTestMocks{
 		Runtime:          rt,
 		ConfigHandler:    mockConfigHandler,
 		Shell:            mockShell,
@@ -204,6 +190,13 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 		ContainerRuntime: mockContainerRuntime,
 		SSHClient:        mockSSHClient,
 	}
+
+	// Apply any overrides
+	for _, opt := range opts {
+		opt(mocks)
+	}
+
+	return mocks
 }
 
 // =============================================================================
@@ -213,7 +206,7 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 func TestNewWorkstation(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 
 		// When
 		workstation, err := NewWorkstation(mocks.Runtime)
@@ -235,7 +228,7 @@ func TestNewWorkstation(t *testing.T) {
 
 	t.Run("NilContext", func(t *testing.T) {
 		// Given
-		_ = setupMocks(t)
+		_ = setupWorkstationMocks(t)
 
 		// When
 		workstation, err := NewWorkstation(nil)
@@ -254,7 +247,7 @@ func TestNewWorkstation(t *testing.T) {
 
 	t.Run("NilConfigHandler", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		rt := &ctxpkg.Runtime{
 			Shell: mocks.Shell,
 		}
@@ -276,7 +269,7 @@ func TestNewWorkstation(t *testing.T) {
 
 	t.Run("NilShell", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		rt := &ctxpkg.Runtime{
 			ConfigHandler: mocks.ConfigHandler,
 		}
@@ -298,7 +291,7 @@ func TestNewWorkstation(t *testing.T) {
 
 	t.Run("NilRuntime", func(t *testing.T) {
 		// Given
-		_ = setupMocks(t)
+		_ = setupWorkstationMocks(t)
 
 		// When
 		workstation, err := NewWorkstation(nil)
@@ -317,7 +310,7 @@ func TestNewWorkstation(t *testing.T) {
 
 	t.Run("CreatesDependencies", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 
 		// When
 		workstation, err := NewWorkstation(mocks.Runtime)
@@ -345,7 +338,7 @@ func TestNewWorkstation(t *testing.T) {
 
 	t.Run("UsesExistingDependencies", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		opts := &Workstation{
 			Runtime:          mocks.Runtime,
 			NetworkManager:   mocks.NetworkManager,
@@ -387,7 +380,7 @@ func TestNewWorkstation(t *testing.T) {
 func TestWorkstation_Up(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime, &Workstation{
 			VirtualMachine:   mocks.VirtualMachine,
 			ContainerRuntime: mocks.ContainerRuntime,
@@ -409,7 +402,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("SetsNoCacheEnvironmentVariable", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime, &Workstation{
 			VirtualMachine:   mocks.VirtualMachine,
 			ContainerRuntime: mocks.ContainerRuntime,
@@ -434,7 +427,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("StartsVirtualMachine", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		vmUpCalled := false
 		mocks.VirtualMachine.UpFunc = func(verbose ...bool) error {
 			vmUpCalled = true
@@ -464,7 +457,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("StartsContainerRuntime", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		containerUpCalled := false
 		mocks.ContainerRuntime.UpFunc = func(verbose ...bool) error {
 			containerUpCalled = true
@@ -494,7 +487,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("ConfiguresNetworking", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		hostRouteCalled := false
 		guestCalled := false
 		dnsCalled := false
@@ -541,7 +534,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("WritesServiceConfigs", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		writeConfigCalled := false
 		for _, service := range mocks.Services {
 			service.WriteConfigFunc = func() error {
@@ -573,7 +566,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("VirtualMachineUpError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		mocks.VirtualMachine.UpFunc = func(verbose ...bool) error {
 			return fmt.Errorf("VM start failed")
 		}
@@ -601,7 +594,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("ContainerRuntimeUpError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		mocks.ContainerRuntime.UpFunc = func(verbose ...bool) error {
 			return fmt.Errorf("container start failed")
 		}
@@ -629,7 +622,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("NetworkConfigurationError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		mocks.NetworkManager.ConfigureHostRouteFunc = func() error {
 			return fmt.Errorf("network config failed")
 		}
@@ -657,7 +650,7 @@ func TestWorkstation_Up(t *testing.T) {
 
 	t.Run("ServiceWriteConfigError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		for _, service := range mocks.Services {
 			service.WriteConfigFunc = func() error {
 				return fmt.Errorf("service config failed")
@@ -689,7 +682,7 @@ func TestWorkstation_Up(t *testing.T) {
 func TestWorkstation_Down(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -706,7 +699,7 @@ func TestWorkstation_Down(t *testing.T) {
 
 	t.Run("StopsContainerRuntime", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		containerDownCalled := false
 		mocks.ContainerRuntime.DownFunc = func() error {
 			containerDownCalled = true
@@ -733,7 +726,7 @@ func TestWorkstation_Down(t *testing.T) {
 
 	t.Run("StopsVirtualMachine", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		vmDownCalled := false
 		mocks.VirtualMachine.DownFunc = func() error {
 			vmDownCalled = true
@@ -760,7 +753,7 @@ func TestWorkstation_Down(t *testing.T) {
 
 	t.Run("ContainerRuntimeDownError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		mocks.ContainerRuntime.DownFunc = func() error {
 			return fmt.Errorf("container stop failed")
 		}
@@ -785,7 +778,7 @@ func TestWorkstation_Down(t *testing.T) {
 
 	t.Run("VirtualMachineDownError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		mocks.VirtualMachine.DownFunc = func() error {
 			return fmt.Errorf("VM stop failed")
 		}
@@ -816,7 +809,7 @@ func TestWorkstation_Down(t *testing.T) {
 func TestWorkstation_createServices(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -839,7 +832,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("DockerDisabled", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		mockConfig := config.NewMockConfigHandler()
 		mockConfig.GetBoolFunc = func(key string, defaultValue ...bool) bool {
 			if key == "docker.enabled" {
@@ -867,7 +860,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("ServiceInitializationError", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -887,7 +880,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("CreatesDNSService", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -907,7 +900,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("CreatesGitLivereloadService", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -927,7 +920,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("CreatesLocalstackService", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -947,7 +940,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("CreatesRegistryServices", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -967,7 +960,7 @@ func TestWorkstation_createServices(t *testing.T) {
 
 	t.Run("CreatesTalosServices", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime)
 		if err != nil {
 			t.Fatalf("Failed to create workstation: %v", err)
@@ -993,7 +986,7 @@ func TestWorkstation_createServices(t *testing.T) {
 func TestWorkstation_Integration(t *testing.T) {
 	t.Run("FullUpDownCycle", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime, &Workstation{
 			VirtualMachine:   mocks.VirtualMachine,
 			ContainerRuntime: mocks.ContainerRuntime,
@@ -1023,7 +1016,7 @@ func TestWorkstation_Integration(t *testing.T) {
 
 	t.Run("MultipleUpDownCycles", func(t *testing.T) {
 		// Given
-		mocks := setupMocks(t)
+		mocks := setupWorkstationMocks(t)
 		workstation, err := NewWorkstation(mocks.Runtime, &Workstation{
 			VirtualMachine:   mocks.VirtualMachine,
 			ContainerRuntime: mocks.ContainerRuntime,

From d05af733b8d7a6318d5b236ffa361dd27706e4fa Mon Sep 17 00:00:00 2001
From: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
Date: Sat, 15 Nov 2025 17:12:16 -0500
Subject: [PATCH 2/4] Standardize network tests

Signed-off-by: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
---
 .../network/colima_network_test.go            | 12 ++--
 .../network/darwin_network_test.go            |  8 +--
 pkg/workstation/network/linux_network_test.go |  8 +--
 pkg/workstation/network/network_test.go       | 55 +++++++------------
 .../network/windows_network_test.go           |  8 +--
 5 files changed, 39 insertions(+), 52 deletions(-)

diff --git a/pkg/workstation/network/colima_network_test.go b/pkg/workstation/network/colima_network_test.go
index fc6df2bc4..55c440cc6 100644
--- a/pkg/workstation/network/colima_network_test.go
+++ b/pkg/workstation/network/colima_network_test.go
@@ -14,9 +14,9 @@ import (
 // =============================================================================
 
 func TestColimaNetworkManager_AssignIPs(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*ColimaNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewColimaNetworkManager(mocks.Runtime, mocks.SSHClient, mocks.SecureShell, mocks.NetworkInterfaceProvider)
 		manager.shims = mocks.Shims
 		return manager, mocks
@@ -42,9 +42,9 @@ func TestColimaNetworkManager_AssignIPs(t *testing.T) {
 }
 
 func TestColimaNetworkManager_ConfigureGuest(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*ColimaNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewColimaNetworkManager(mocks.Runtime, mocks.SSHClient, mocks.SecureShell, mocks.NetworkInterfaceProvider)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})
@@ -300,9 +300,9 @@ func TestColimaNetworkManager_ConfigureGuest(t *testing.T) {
 }
 
 func TestColimaNetworkManager_getHostIP(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*ColimaNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewColimaNetworkManager(mocks.Runtime, mocks.SSHClient, mocks.SecureShell, mocks.NetworkInterfaceProvider)
 		manager.AssignIPs([]services.Service{})
 		return manager, mocks
diff --git a/pkg/workstation/network/darwin_network_test.go b/pkg/workstation/network/darwin_network_test.go
index 32dd8842e..2dd01a5df 100644
--- a/pkg/workstation/network/darwin_network_test.go
+++ b/pkg/workstation/network/darwin_network_test.go
@@ -16,9 +16,9 @@ import (
 // =============================================================================
 
 func TestDarwinNetworkManager_ConfigureHostRoute(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})
@@ -164,9 +164,9 @@ func TestDarwinNetworkManager_ConfigureHostRoute(t *testing.T) {
 }
 
 func TestDarwinNetworkManager_ConfigureDNS(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})
diff --git a/pkg/workstation/network/linux_network_test.go b/pkg/workstation/network/linux_network_test.go
index 4c8ef8fa3..29535941f 100644
--- a/pkg/workstation/network/linux_network_test.go
+++ b/pkg/workstation/network/linux_network_test.go
@@ -17,9 +17,9 @@ import (
 // =============================================================================
 
 func TestLinuxNetworkManager_ConfigureHostRoute(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})
@@ -166,9 +166,9 @@ func TestLinuxNetworkManager_ConfigureHostRoute(t *testing.T) {
 }
 
 func TestLinuxNetworkManager_ConfigureDNS(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})
diff --git a/pkg/workstation/network/network_test.go b/pkg/workstation/network/network_test.go
index 87f622d81..ab5f555b1 100644
--- a/pkg/workstation/network/network_test.go
+++ b/pkg/workstation/network/network_test.go
@@ -18,7 +18,7 @@ import (
 // Test Setup
 // =============================================================================
 
-type Mocks struct {
+type NetworkTestMocks struct {
 	Runtime                  *runtime.Runtime
 	ConfigHandler            config.ConfigHandler
 	Shell                    *shell.MockShell
@@ -29,14 +29,7 @@ type Mocks struct {
 	Shims                    *Shims
 }
 
-type SetupOptions struct {
-	ConfigHandler config.ConfigHandler
-	ConfigStr     string
-}
-
-func setupShims(t *testing.T) *Shims {
-	t.Helper()
-
+func setupDefaultShims() *Shims {
 	return &Shims{
 		Stat:      func(path string) (os.FileInfo, error) { return nil, nil },
 		WriteFile: func(path string, data []byte, perm os.FileMode) error { return nil },
@@ -46,7 +39,7 @@ func setupShims(t *testing.T) *Shims {
 	}
 }
 
-func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupNetworkMocks(t *testing.T, opts ...func(*NetworkTestMocks)) *NetworkTestMocks {
 	t.Helper()
 
 	// Store original directory and create temp dir
@@ -77,13 +70,8 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 		return tmpDir, nil
 	}
 
-	// Create config handler if not provided
-	var configHandler config.ConfigHandler
-	if len(opts) > 0 && opts[0].ConfigHandler != nil {
-		configHandler = opts[0].ConfigHandler
-	} else {
-		configHandler = config.NewConfigHandler(mockShell)
-	}
+	// Create config handler
+	configHandler := config.NewConfigHandler(mockShell)
 
 	configYAML := `
 version: v1alpha1
@@ -101,13 +89,6 @@ contexts:
 		t.Fatalf("Failed to load config: %v", err)
 	}
 
-	// Load optional config if provided
-	if len(opts) > 0 && opts[0].ConfigStr != "" {
-		if err := configHandler.LoadConfigString(opts[0].ConfigStr); err != nil {
-			t.Fatalf("Failed to load config string: %v", err)
-		}
-	}
-
 	// Configure mock shell functions
 	mockShell.ExecFunc = func(command string, args ...string) (string, error) {
 		return "", nil
@@ -192,7 +173,7 @@ contexts:
 	}
 
 	// Create mocks struct with references to the same instances
-	mocks := &Mocks{
+	mocks := &NetworkTestMocks{
 		Runtime:                  rt,
 		ConfigHandler:            configHandler,
 		Shell:                    mockShell,
@@ -200,11 +181,16 @@ contexts:
 		SSHClient:                mockSSHClient,
 		NetworkInterfaceProvider: mockNetworkInterfaceProvider,
 		Services:                 []*services.MockService{mockService1, mockService2},
-		Shims:                    setupShims(t),
+		Shims:                    setupDefaultShims(),
 	}
 
 	configHandler.SetContext("mock-context")
 
+	// Apply any overrides
+	for _, opt := range opts {
+		opt(mocks)
+	}
+
 	return mocks
 }
 
@@ -235,9 +221,9 @@ func TestNetworkManager_NewNetworkManager(t *testing.T) {
 // =============================================================================
 
 func TestNetworkManager_AssignIPs(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		return manager, mocks
@@ -358,8 +344,9 @@ func TestNetworkManager_AssignIPs(t *testing.T) {
 		}
 
 		// Setup with mock config handler
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupNetworkMocks(t, func(m *NetworkTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
@@ -408,9 +395,9 @@ func TestNetworkManager_AssignIPs(t *testing.T) {
 }
 
 func TestNetworkManager_ConfigureGuest(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		return manager, mocks
@@ -431,9 +418,9 @@ func TestNetworkManager_ConfigureGuest(t *testing.T) {
 }
 
 func TestNetworkManager_assignIPAddresses(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		return manager, mocks
diff --git a/pkg/workstation/network/windows_network_test.go b/pkg/workstation/network/windows_network_test.go
index 6b45d8f48..00da3b1b4 100644
--- a/pkg/workstation/network/windows_network_test.go
+++ b/pkg/workstation/network/windows_network_test.go
@@ -16,9 +16,9 @@ import (
 // =============================================================================
 
 func TestWindowsNetworkManager_ConfigureHostRoute(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})
@@ -139,9 +139,9 @@ func TestWindowsNetworkManager_ConfigureHostRoute(t *testing.T) {
 }
 
 func TestWindowsNetworkManager_ConfigureDNS(t *testing.T) {
-	setup := func(t *testing.T) (*BaseNetworkManager, *Mocks) {
+	setup := func(t *testing.T) (*BaseNetworkManager, *NetworkTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupNetworkMocks(t)
 		manager := NewBaseNetworkManager(mocks.Runtime)
 		manager.shims = mocks.Shims
 		manager.AssignIPs([]services.Service{})

From e6c369c28971d06298e3ca96e7c1d638a7dde8d5 Mon Sep 17 00:00:00 2001
From: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
Date: Sat, 15 Nov 2025 17:18:12 -0500
Subject: [PATCH 3/4] Standardize workstation/service tests

Signed-off-by: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
---
 pkg/workstation/services/dns_service_test.go  | 32 ++++----
 .../services/git_livereload_service_test.go   |  9 +--
 .../services/localstack_service_test.go       | 10 +--
 .../services/registry_service_test.go         | 54 +++++++------
 pkg/workstation/services/service_test.go      | 78 ++++++++-----------
 .../services/talos_service_test.go            | 78 +++++++++++++------
 6 files changed, 141 insertions(+), 120 deletions(-)

diff --git a/pkg/workstation/services/dns_service_test.go b/pkg/workstation/services/dns_service_test.go
index f94a61a29..d82d3b345 100644
--- a/pkg/workstation/services/dns_service_test.go
+++ b/pkg/workstation/services/dns_service_test.go
@@ -15,11 +15,11 @@ import (
 // =============================================================================
 
 // setupDnsMocks creates and returns mock components for DNS service tests
-func setupDnsMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupDnsMocks(t *testing.T, opts ...func(*ServicesTestMocks)) *ServicesTestMocks {
 	t.Helper()
 
-	// Create base mocks using setupMocks
-	mocks := setupMocks(t, opts...)
+	// Create base mocks using setupServicesMocks
+	mocks := setupServicesMocks(t, opts...)
 
 	// Set up shell project root
 	mocks.Shell.GetProjectRootFunc = func() (string, error) {
@@ -34,7 +34,7 @@ func setupDnsMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 // =============================================================================
 
 func TestNewDNSService(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -59,7 +59,7 @@ func TestNewDNSService(t *testing.T) {
 // =============================================================================
 
 func TestDNSService_Initialize(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -80,7 +80,7 @@ func TestDNSService_Initialize(t *testing.T) {
 }
 
 func TestDNSService_SetAddress(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -114,8 +114,9 @@ func TestDNSService_SetAddress(t *testing.T) {
 		mockConfigHandler.SetFunc = func(key string, value any) error {
 			return fmt.Errorf("mocked error setting address")
 		}
-		mocks := setupDnsMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupDnsMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 		service := NewDNSService(mocks.Runtime)
 		service.shims = mocks.Shims
@@ -136,7 +137,7 @@ func TestDNSService_SetAddress(t *testing.T) {
 }
 
 func TestDNSService_GetComposeConfig(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -216,7 +217,7 @@ func TestDNSService_GetComposeConfig(t *testing.T) {
 }
 
 func TestDNSService_WriteConfig(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -595,7 +596,6 @@ func TestDNSService_WriteConfig(t *testing.T) {
 		}
 	})
 
-
 	t.Run("SuccessRemovingCorefileDirectory", func(t *testing.T) {
 		// Given a DNSService with mock components
 		service, mocks := setup(t)
@@ -642,7 +642,7 @@ func TestDNSService_WriteConfig(t *testing.T) {
 }
 
 func TestDNSService_SetName(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -667,7 +667,7 @@ func TestDNSService_SetName(t *testing.T) {
 }
 
 func TestDNSService_GetName(t *testing.T) {
-	setupSuccess := func(t *testing.T) (*DNSService, *Mocks) {
+	setupSuccess := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -677,7 +677,7 @@ func TestDNSService_GetName(t *testing.T) {
 		return service, mocks
 	}
 
-	setupError := func(t *testing.T) (*DNSService, *Mocks) {
+	setupError := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -720,7 +720,7 @@ func TestDNSService_GetName(t *testing.T) {
 }
 
 func TestDNSService_GetHostname(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
@@ -764,7 +764,7 @@ func TestDNSService_GetHostname(t *testing.T) {
 }
 
 func TestDNSService_SupportsWildcard(t *testing.T) {
-	setup := func(t *testing.T) (*DNSService, *Mocks) {
+	setup := func(t *testing.T) (*DNSService, *ServicesTestMocks) {
 		t.Helper()
 		mocks := setupDnsMocks(t)
 		service := NewDNSService(mocks.Runtime)
diff --git a/pkg/workstation/services/git_livereload_service_test.go b/pkg/workstation/services/git_livereload_service_test.go
index 2579ddb27..a855a68ac 100644
--- a/pkg/workstation/services/git_livereload_service_test.go
+++ b/pkg/workstation/services/git_livereload_service_test.go
@@ -18,7 +18,7 @@ import (
 func TestGitLivereloadService_NewGitLivereloadService(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		// Given a set of mock components
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 
 		// When a new GitLivereloadService is created
 		gitLivereloadService := NewGitLivereloadService(mocks.Runtime)
@@ -34,7 +34,7 @@ func TestGitLivereloadService_NewGitLivereloadService(t *testing.T) {
 func TestGitLivereloadService_GetComposeConfig(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		// Given a mock config handler, shell, context, and service
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		gitLivereloadService := NewGitLivereloadService(mocks.Runtime)
 
 		// Set the service name
@@ -63,10 +63,9 @@ func TestGitLivereloadService_GetComposeConfig(t *testing.T) {
 		}
 	})
 
-
 	t.Run("SuccessWithRsyncInclude", func(t *testing.T) {
 		// Given a mock config handler, shell, context, and service with rsync_include configured
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		gitLivereloadService := NewGitLivereloadService(mocks.Runtime)
 
 		// Set the service name
@@ -109,7 +108,7 @@ func TestGitLivereloadService_GetComposeConfig(t *testing.T) {
 
 	t.Run("SuccessWithoutRsyncInclude", func(t *testing.T) {
 		// Given a mock config handler, shell, context, and service without rsync_include configured
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		gitLivereloadService := NewGitLivereloadService(mocks.Runtime)
 
 		// Set the service name
diff --git a/pkg/workstation/services/localstack_service_test.go b/pkg/workstation/services/localstack_service_test.go
index 64a07e32b..f7f5dea89 100644
--- a/pkg/workstation/services/localstack_service_test.go
+++ b/pkg/workstation/services/localstack_service_test.go
@@ -10,9 +10,9 @@ import (
 
 // TestLocalstackService_GetComposeConfig tests the GetComposeConfig method
 func TestLocalstackService_GetComposeConfig(t *testing.T) {
-	setup := func(t *testing.T) (*LocalstackService, *Mocks) {
+	setup := func(t *testing.T) (*LocalstackService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		service := NewLocalstackService(mocks.Runtime)
 		service.shims = mocks.Shims
 		service.SetName("aws")
@@ -34,7 +34,6 @@ func TestLocalstackService_GetComposeConfig(t *testing.T) {
 			t.Fatalf("failed to set localstack services: %v", err)
 		}
 
-
 		// When: GetComposeConfig is called
 		composeConfig, err := service.GetComposeConfig()
 		if err != nil {
@@ -75,7 +74,6 @@ func TestLocalstackService_GetComposeConfig(t *testing.T) {
 			t.Fatalf("failed to set localstack services: %v", err)
 		}
 
-
 		// When: GetComposeConfig is called
 		composeConfig, err := service.GetComposeConfig()
 		if err != nil {
@@ -99,9 +97,9 @@ func TestLocalstackService_GetComposeConfig(t *testing.T) {
 
 // TestLocalstackService_SupportsWildcard tests the SupportsWildcard method
 func TestLocalstackService_SupportsWildcard(t *testing.T) {
-	setup := func(t *testing.T) (*LocalstackService, *Mocks) {
+	setup := func(t *testing.T) (*LocalstackService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		service := NewLocalstackService(mocks.Runtime)
 		service.shims = mocks.Shims
 		service.SetName("aws")
diff --git a/pkg/workstation/services/registry_service_test.go b/pkg/workstation/services/registry_service_test.go
index 945dde2e7..45ae668f4 100644
--- a/pkg/workstation/services/registry_service_test.go
+++ b/pkg/workstation/services/registry_service_test.go
@@ -18,9 +18,9 @@ import (
 // =============================================================================
 
 func TestRegistryService_NewRegistryService(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		service := NewRegistryService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -42,9 +42,9 @@ func TestRegistryService_NewRegistryService(t *testing.T) {
 // =============================================================================
 
 func TestRegistryService_GetComposeConfig(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		service := NewRegistryService(mocks.Runtime)
 		service.shims = mocks.Shims
 		service.SetName("registry")
@@ -169,9 +169,9 @@ func TestRegistryService_GetComposeConfig(t *testing.T) {
 }
 
 func TestRegistryService_SetAddress(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		// Load initial config
 		configYAML := `
 version: v1alpha1
@@ -374,8 +374,9 @@ contexts:
 	t.Run("BaseServiceError", func(t *testing.T) {
 		// Given a registry service with mock components
 		mockConfigHandler := config.NewMockConfigHandler()
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 
 		service := NewRegistryService(mocks.Runtime)
@@ -397,8 +398,9 @@ contexts:
 	t.Run("ErrorSettingHostname", func(t *testing.T) {
 		// Given a registry service with mock components
 		mockConfigHandler := config.NewMockConfigHandler()
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 
 		service := NewRegistryService(mocks.Runtime)
@@ -425,8 +427,9 @@ contexts:
 	t.Run("ErrorSettingHostPort", func(t *testing.T) {
 		// Given a registry service with mock components
 		mockConfigHandler := config.NewMockConfigHandler()
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 
 		service := NewRegistryService(mocks.Runtime)
@@ -476,8 +479,9 @@ contexts:
 	t.Run("ErrorSettingRegistryURL", func(t *testing.T) {
 		// Given a registry service with mock components
 		mockConfigHandler := config.NewMockConfigHandler()
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 
 		service := NewRegistryService(mocks.Runtime)
@@ -530,7 +534,7 @@ contexts:
 }
 
 func TestRegistryService_GetHostname(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		mockConfigHandler := config.NewMockConfigHandler()
 		mockConfigHandler.GetStringFunc = func(key string, defaultValue ...string) string {
 			if key == "dns.domain" {
@@ -538,8 +542,9 @@ func TestRegistryService_GetHostname(t *testing.T) {
 			}
 			return defaultValue[0]
 		}
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 		service := NewRegistryService(mocks.Runtime)
 		return service, mocks
@@ -606,7 +611,7 @@ func TestRegistryService_GetHostname(t *testing.T) {
 }
 
 func TestRegistryService_GetContainerName(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		mockConfigHandler := config.NewMockConfigHandler()
 		mockConfigHandler.GetStringFunc = func(key string, defaultValue ...string) string {
 			if key == "dns.domain" {
@@ -614,8 +619,9 @@ func TestRegistryService_GetContainerName(t *testing.T) {
 			}
 			return defaultValue[0]
 		}
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 		service := NewRegistryService(mocks.Runtime)
 		return service, mocks
@@ -638,9 +644,9 @@ func TestRegistryService_GetContainerName(t *testing.T) {
 }
 
 func TestRegistryService_GetName(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		service := NewRegistryService(mocks.Runtime)
 		service.shims = mocks.Shims
 		service.SetName("registry")
@@ -658,9 +664,9 @@ func TestRegistryService_GetName(t *testing.T) {
 }
 
 func TestRegistryService_SupportsWildcard(t *testing.T) {
-	setup := func(t *testing.T) (*RegistryService, *Mocks) {
+	setup := func(t *testing.T) (*RegistryService, *ServicesTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupServicesMocks(t)
 		service := NewRegistryService(mocks.Runtime)
 		service.shims = mocks.Shims
 		service.SetName("registry")
diff --git a/pkg/workstation/services/service_test.go b/pkg/workstation/services/service_test.go
index 4ac99cd13..5f53ded11 100644
--- a/pkg/workstation/services/service_test.go
+++ b/pkg/workstation/services/service_test.go
@@ -31,20 +31,14 @@ func (m *mockFileInfo) ModTime() time.Time { return time.Now() }
 func (m *mockFileInfo) IsDir() bool        { return m.isDir }
 func (m *mockFileInfo) Sys() interface{}   { return nil }
 
-type Mocks struct {
+type ServicesTestMocks struct {
 	Runtime       *runtime.Runtime
 	ConfigHandler config.ConfigHandler
 	Shell         *shell.MockShell
 	Shims         *Shims
 }
 
-type SetupOptions struct {
-	ConfigHandler config.ConfigHandler
-	ConfigStr     string
-}
-
-func setupShims(t *testing.T) *Shims {
-	t.Helper()
+func setupDefaultShims() *Shims {
 	shims := NewShims()
 
 	shims.Getwd = func() (string, error) {
@@ -88,7 +82,7 @@ func setupShims(t *testing.T) *Shims {
 	return shims
 }
 
-func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupServicesMocks(t *testing.T, opts ...func(*ServicesTestMocks)) *ServicesTestMocks {
 	t.Helper()
 
 	// Store original directory and create temp dir
@@ -120,13 +114,8 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 		return tmpDir, nil
 	}
 
-	// Create config handler if not provided
-	var configHandler config.ConfigHandler
-	if len(opts) > 0 && opts[0].ConfigHandler != nil {
-		configHandler = opts[0].ConfigHandler
-	} else {
-		configHandler = config.NewConfigHandler(mockShell)
-	}
+	// Create config handler
+	configHandler := config.NewConfigHandler(mockShell)
 
 	configHandler.SetContext("mock-context")
 
@@ -152,13 +141,6 @@ contexts:
 		t.Fatalf("Failed to load config: %v", err)
 	}
 
-	// Load optional config if provided
-	if len(opts) > 0 && opts[0].ConfigStr != "" {
-		if err := configHandler.LoadConfigString(opts[0].ConfigStr); err != nil {
-			t.Fatalf("Failed to load config string: %v", err)
-		}
-	}
-
 	rt := &runtime.Runtime{
 		ProjectRoot:   tmpDir,
 		ConfigRoot:    tmpDir,
@@ -168,12 +150,19 @@ contexts:
 		Shell:         mockShell,
 	}
 
-	return &Mocks{
+	mocks := &ServicesTestMocks{
 		Runtime:       rt,
 		ConfigHandler: configHandler,
 		Shell:         mockShell,
-		Shims:         setupShims(t),
+		Shims:         setupDefaultShims(),
 	}
+
+	// Apply any overrides
+	for _, opt := range opts {
+		opt(mocks)
+	}
+
+	return mocks
 }
 
 // =============================================================================
@@ -181,8 +170,8 @@ contexts:
 // =============================================================================
 
 func TestBaseService_NewBaseService(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		return service, mocks
 	}
@@ -209,8 +198,8 @@ func TestBaseService_NewBaseService(t *testing.T) {
 }
 
 func TestBaseService_WriteConfig(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -231,8 +220,8 @@ func TestBaseService_WriteConfig(t *testing.T) {
 }
 
 func TestBaseService_SetAddress(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -278,8 +267,8 @@ func TestBaseService_SetAddress(t *testing.T) {
 }
 
 func TestBaseService_GetAddress(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -302,8 +291,8 @@ func TestBaseService_GetAddress(t *testing.T) {
 }
 
 func TestBaseService_GetName(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -326,8 +315,8 @@ func TestBaseService_GetName(t *testing.T) {
 }
 
 func TestBaseService_GetHostname(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -366,8 +355,8 @@ func TestBaseService_GetHostname(t *testing.T) {
 }
 
 func TestBaseService_IsLocalhostMode(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -407,8 +396,8 @@ func TestBaseService_IsLocalhostMode(t *testing.T) {
 }
 
 func TestBaseService_SupportsWildcard(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
-		mocks := setupMocks(t)
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
+		mocks := setupServicesMocks(t)
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
 		return service, mocks
@@ -429,7 +418,7 @@ func TestBaseService_SupportsWildcard(t *testing.T) {
 }
 
 func TestBaseService_GetContainerName(t *testing.T) {
-	setup := func(t *testing.T) (*BaseService, *Mocks) {
+	setup := func(t *testing.T) (*BaseService, *ServicesTestMocks) {
 		mockConfigHandler := config.NewMockConfigHandler()
 		mockConfigHandler.GetStringFunc = func(key string, defaultValue ...string) string {
 			if key == "dns.domain" {
@@ -437,8 +426,9 @@ func TestBaseService_GetContainerName(t *testing.T) {
 			}
 			return defaultValue[0]
 		}
-		mocks := setupMocks(t, &SetupOptions{
-			ConfigHandler: mockConfigHandler,
+		mocks := setupServicesMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = mockConfigHandler
+			m.Runtime.ConfigHandler = mockConfigHandler
 		})
 		service := NewBaseService(mocks.Runtime)
 		service.shims = mocks.Shims
diff --git a/pkg/workstation/services/talos_service_test.go b/pkg/workstation/services/talos_service_test.go
index d47068cff..e90ad2ddc 100644
--- a/pkg/workstation/services/talos_service_test.go
+++ b/pkg/workstation/services/talos_service_test.go
@@ -8,6 +8,8 @@ import (
 	"testing"
 
 	"github.com/windsorcli/cli/pkg/constants"
+	"github.com/windsorcli/cli/pkg/runtime/config"
+	"github.com/windsorcli/cli/pkg/runtime/shell"
 )
 
 // =============================================================================
@@ -15,18 +17,22 @@ import (
 // =============================================================================
 
 // setupTalosServiceMocks creates and returns mock components for TalosService tests
-func setupTalosServiceMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupTalosServiceMocks(t *testing.T, opts ...func(*ServicesTestMocks)) *ServicesTestMocks {
 	t.Helper()
 
-	// Create base mocks using setupMocks
-	mocks := setupMocks(t, opts...)
+	// Create base mocks using setupServicesMocks
+	mocks := setupServicesMocks(t, opts...)
 
-	// Load config - use provided config if available, otherwise use default
-	var configToLoad string
-	if len(opts) > 0 && opts[0].ConfigStr != "" {
-		configToLoad = opts[0].ConfigStr
-	} else {
-		configToLoad = fmt.Sprintf(`
+	// Check if config handler was overridden (compare with the default one from setupServicesMocks)
+	// We'll check by seeing if the config handler is different from what setupServicesMocks created
+	// For simplicity, we'll just check if opts were provided - if they were, assume custom config might be set
+	// and skip loading default talos config. Tests that need default config won't provide opts.
+	hasCustomConfig := len(opts) > 0
+
+	// Only load default talos config if no custom config handler was provided
+	if !hasCustomConfig {
+		// Load config
+		configToLoad := fmt.Sprintf(`
 version: v1alpha1
 contexts:
   mock-context:
@@ -71,10 +77,10 @@ contexts:
 `, constants.DefaultTalosAPIPort,
 			constants.DefaultTalosWorkerCPU,
 			constants.DefaultTalosWorkerRAM)
-	}
 
-	if err := mocks.ConfigHandler.LoadConfigString(configToLoad); err != nil {
-		t.Fatalf("Failed to load config: %v", err)
+		if err := mocks.ConfigHandler.LoadConfigString(configToLoad); err != nil {
+			t.Fatalf("Failed to load config: %v", err)
+		}
 	}
 
 	mocks.Shell.GetProjectRootFunc = func() (string, error) {
@@ -123,7 +129,7 @@ func TestTalosService_NewTalosService(t *testing.T) {
 
 // TestTalosService_SetAddress tests the SetAddress method of TalosService
 func TestTalosService_SetAddress(t *testing.T) {
-	setup := func(t *testing.T) (*TalosService, *Mocks) {
+	setup := func(t *testing.T) (*TalosService, *ServicesTestMocks) {
 		t.Helper()
 
 		// Reset package-level variables
@@ -829,7 +835,7 @@ func TestTalosService_SetAddress(t *testing.T) {
 
 // TestTalosService_GetComposeConfig tests the GetComposeConfig method of TalosService
 func TestTalosService_GetComposeConfig(t *testing.T) {
-	setup := func(t *testing.T) (*TalosService, *Mocks) {
+	setup := func(t *testing.T) (*TalosService, *ServicesTestMocks) {
 		t.Helper()
 
 		// Reset package-level variables
@@ -848,7 +854,7 @@ func TestTalosService_GetComposeConfig(t *testing.T) {
 		return service, mocks
 	}
 
-	setupWorker := func(t *testing.T) (*TalosService, *Mocks) {
+	setupWorker := func(t *testing.T) (*TalosService, *ServicesTestMocks) {
 		t.Helper()
 
 		// Reset package-level variables
@@ -1049,8 +1055,14 @@ func TestTalosService_GetComposeConfig(t *testing.T) {
 
 	t.Run("SuccessEmptyConfig", func(t *testing.T) {
 		// Given a TalosService with mock components and empty cluster config
-		emptyConfig := &SetupOptions{
-			ConfigStr: `
+		// Create a new config handler with empty config
+		mockShell := shell.NewMockShell()
+		mockShell.GetProjectRootFunc = func() (string, error) {
+			return "/mock/project/root", nil
+		}
+		configHandler := config.NewConfigHandler(mockShell)
+		configHandler.SetContext("mock-context")
+		configStr := `
 version: v1alpha1
 contexts:
   mock-context:
@@ -1058,9 +1070,14 @@ contexts:
       domain: test
     vm:
       driver: docker-desktop
-`,
-		}
-		mocks := setupTalosServiceMocks(t, emptyConfig)
+`
+		if err := configHandler.LoadConfigString(configStr); err != nil {
+			t.Fatalf("Failed to load config: %v", err)
+		}
+		mocks := setupTalosServiceMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = configHandler
+			m.Runtime.ConfigHandler = configHandler
+		})
 		service := NewTalosService(mocks.Runtime, "controlplane")
 		service.shims = mocks.Shims
 		service.SetName("controlplane1")
@@ -1084,8 +1101,14 @@ contexts:
 
 	t.Run("EmptyConfig", func(t *testing.T) {
 		// Given a TalosService with mock components and empty cluster config
-		emptyConfig := &SetupOptions{
-			ConfigStr: `
+		// Create a new config handler with empty config
+		mockShell := shell.NewMockShell()
+		mockShell.GetProjectRootFunc = func() (string, error) {
+			return "/mock/project/root", nil
+		}
+		configHandler := config.NewConfigHandler(mockShell)
+		configHandler.SetContext("mock-context")
+		configStr := `
 version: v1alpha1
 contexts:
   mock-context:
@@ -1093,9 +1116,14 @@ contexts:
       domain: test
     vm:
       driver: docker-desktop
-`,
-		}
-		mocks := setupTalosServiceMocks(t, emptyConfig)
+`
+		if err := configHandler.LoadConfigString(configStr); err != nil {
+			t.Fatalf("Failed to load config: %v", err)
+		}
+		mocks := setupTalosServiceMocks(t, func(m *ServicesTestMocks) {
+			m.ConfigHandler = configHandler
+			m.Runtime.ConfigHandler = configHandler
+		})
 		service := NewTalosService(mocks.Runtime, "controlplane")
 		service.shims = mocks.Shims
 		service.SetName("controlplane1")

From ca7f74e044a53ab58130d1c8aaa4d1523f4337c6 Mon Sep 17 00:00:00 2001
From: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
Date: Sat, 15 Nov 2025 17:23:11 -0500
Subject: [PATCH 4/4] Enhance virt tests

Signed-off-by: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com>
---
 pkg/workstation/virt/colima_virt_test.go | 25 +++++------
 pkg/workstation/virt/docker_virt_test.go | 22 ++++------
 pkg/workstation/virt/virt_test.go        | 56 ++++++++----------------
 3 files changed, 36 insertions(+), 67 deletions(-)

diff --git a/pkg/workstation/virt/colima_virt_test.go b/pkg/workstation/virt/colima_virt_test.go
index be070d804..0b24fd532 100644
--- a/pkg/workstation/virt/colima_virt_test.go
+++ b/pkg/workstation/virt/colima_virt_test.go
@@ -24,16 +24,11 @@ import (
 // Test Setup
 // =============================================================================
 
-func setupColimaMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupColimaMocks(t *testing.T, opts ...func(*VirtTestMocks)) *VirtTestMocks {
 	t.Helper()
 
-	var options *SetupOptions
-	if len(opts) > 0 {
-		options = opts[0]
-	}
-
 	// Set up mocks and shell
-	mocks := setupMocks(t, options)
+	mocks := setupVirtMocks(t, opts...)
 
 	// Set up shell mock for GetVMInfo
 	mocks.Shell.ExecSilentFunc = func(command string, args ...string) (string, error) {
@@ -126,7 +121,7 @@ contexts:
 // =============================================================================
 
 func TestColimaVirt_Initialize(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupColimaMocks(t)
 		colimaVirt := NewColimaVirt(mocks.Runtime)
@@ -164,7 +159,7 @@ func TestColimaVirt_Initialize(t *testing.T) {
 }
 
 func TestColimaVirt_WriteConfig(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupColimaMocks(t)
 
@@ -466,7 +461,7 @@ func TestColimaVirt_WriteConfig(t *testing.T) {
 }
 
 func TestColimaVirt_Up(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupColimaMocks(t)
 		colimaVirt := NewColimaVirt(mocks.Runtime)
@@ -582,7 +577,7 @@ func TestColimaVirt_Up(t *testing.T) {
 }
 
 func TestColimaVirt_Down(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupColimaMocks(t)
 		colimaVirt := NewColimaVirt(mocks.Runtime)
@@ -678,9 +673,9 @@ func TestColimaVirt_Down(t *testing.T) {
 
 // TestColimaVirt_getArch tests the getArch method of the ColimaVirt component.
 func TestColimaVirt_getArch(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupVirtMocks(t)
 		colimaVirt := NewColimaVirt(mocks.Runtime)
 		colimaVirt.shims = mocks.Shims
 		return colimaVirt, mocks
@@ -731,7 +726,7 @@ func TestColimaVirt_getArch(t *testing.T) {
 
 // TestColimaVirt_getDefaultValues tests the getDefaultValues method of the ColimaVirt component.
 func TestColimaVirt_getDefaultValues(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupColimaMocks(t)
 		colimaVirt := NewColimaVirt(mocks.Runtime)
@@ -802,7 +797,7 @@ func TestColimaVirt_getDefaultValues(t *testing.T) {
 
 // TestColimaVirt_startColima tests the startColima method of the ColimaVirt component.
 func TestColimaVirt_startColima(t *testing.T) {
-	setup := func(t *testing.T) (*ColimaVirt, *Mocks) {
+	setup := func(t *testing.T) (*ColimaVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupColimaMocks(t)
 		colimaVirt := NewColimaVirt(mocks.Runtime)
diff --git a/pkg/workstation/virt/docker_virt_test.go b/pkg/workstation/virt/docker_virt_test.go
index eb284095f..b04fefc67 100644
--- a/pkg/workstation/virt/docker_virt_test.go
+++ b/pkg/workstation/virt/docker_virt_test.go
@@ -19,17 +19,11 @@ import (
 // Test Setup
 // =============================================================================
 
-func setupDockerMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupDockerMocks(t *testing.T, opts ...func(*VirtTestMocks)) *VirtTestMocks {
 	t.Helper()
 
-	// Process options with defaults
-	options := &SetupOptions{}
-	if len(opts) > 0 && opts[0] != nil {
-		options = opts[0]
-	}
-
 	// Set up base mocks
-	mocks := setupMocks(t, options)
+	mocks := setupVirtMocks(t, opts...)
 
 	// Load Docker-specific config
 	configStr := `
@@ -162,7 +156,7 @@ contexts:
 
 // TestDockerVirt_Initialize tests the initialization of the DockerVirt component.
 func TestDockerVirt_Initialize(t *testing.T) {
-	setup := func(t *testing.T) (*DockerVirt, *Mocks) {
+	setup := func(t *testing.T) (*DockerVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupDockerMocks(t)
 		dockerVirt := NewDockerVirt(mocks.Runtime, []services.Service{})
@@ -301,7 +295,7 @@ func TestDockerVirt_Initialize(t *testing.T) {
 
 // TestDockerVirt_Up tests the Up method of the DockerVirt component.
 func TestDockerVirt_Up(t *testing.T) {
-	setup := func(t *testing.T) (*DockerVirt, *Mocks) {
+	setup := func(t *testing.T) (*DockerVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupDockerMocks(t)
 		dockerVirt := NewDockerVirt(mocks.Runtime, []services.Service{})
@@ -501,7 +495,7 @@ func TestDockerVirt_Up(t *testing.T) {
 
 // TestDockerVirt_Down tests the Down method of the DockerVirt component.
 func TestDockerVirt_Down(t *testing.T) {
-	setup := func(t *testing.T) (*DockerVirt, *Mocks) {
+	setup := func(t *testing.T) (*DockerVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupDockerMocks(t)
 		dockerVirt := NewDockerVirt(mocks.Runtime, []services.Service{})
@@ -663,7 +657,7 @@ func TestDockerVirt_Down(t *testing.T) {
 
 // TestDockerVirt_WriteConfig tests the WriteConfig method of the DockerVirt component.
 func TestDockerVirt_WriteConfig(t *testing.T) {
-	setup := func(t *testing.T) (*DockerVirt, *Mocks) {
+	setup := func(t *testing.T) (*DockerVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupDockerMocks(t)
 		dockerVirt := NewDockerVirt(mocks.Runtime, []services.Service{})
@@ -764,7 +758,7 @@ func TestDockerVirt_WriteConfig(t *testing.T) {
 
 // TestDockerVirt_DetermineComposeCommand tests the determineComposeCommand method of the DockerVirt component.
 func TestDockerVirt_DetermineComposeCommand(t *testing.T) {
-	setup := func(t *testing.T) (*DockerVirt, *Mocks) {
+	setup := func(t *testing.T) (*DockerVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupDockerMocks(t)
 		dockerVirt := NewDockerVirt(mocks.Runtime, []services.Service{})
@@ -877,7 +871,7 @@ func TestDockerVirt_DetermineComposeCommand(t *testing.T) {
 
 // TestDockerVirt_GetFullComposeConfig tests the getFullComposeConfig method of the DockerVirt component.
 func TestDockerVirt_GetFullComposeConfig(t *testing.T) {
-	setup := func(t *testing.T) (*DockerVirt, *Mocks) {
+	setup := func(t *testing.T) (*DockerVirt, *VirtTestMocks) {
 		t.Helper()
 		mocks := setupDockerMocks(t)
 		dockerVirt := NewDockerVirt(mocks.Runtime, []services.Service{})
diff --git a/pkg/workstation/virt/virt_test.go b/pkg/workstation/virt/virt_test.go
index 7a59acfec..9cadd9a28 100644
--- a/pkg/workstation/virt/virt_test.go
+++ b/pkg/workstation/virt/virt_test.go
@@ -23,7 +23,7 @@ import (
 // Test Setup
 // =============================================================================
 
-type Mocks struct {
+type VirtTestMocks struct {
 	Runtime       *runtime.Runtime
 	ConfigHandler config.ConfigHandler
 	Shell         *shell.MockShell
@@ -31,14 +31,8 @@ type Mocks struct {
 	Service       *services.MockService
 }
 
-type SetupOptions struct {
-	ConfigHandler config.ConfigHandler
-	ConfigStr     string
-}
-
-// setupShims creates a new Shims instance with default implementations
-func setupShims(t *testing.T) *Shims {
-	t.Helper()
+// setupDefaultShims creates a new Shims instance with default implementations
+func setupDefaultShims() *Shims {
 	shims := &Shims{
 		Setenv: func(key, value string) error {
 			return os.Setenv(key, value)
@@ -87,15 +81,10 @@ func setupShims(t *testing.T) *Shims {
 		},
 	}
 
-	t.Cleanup(func() {
-		os.Unsetenv("COMPOSE_FILE")
-		os.Unsetenv("WINDSOR_CONTEXT")
-	})
-
 	return shims
 }
 
-func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
+func setupVirtMocks(t *testing.T, opts ...func(*VirtTestMocks)) *VirtTestMocks {
 	t.Helper()
 
 	// Store original directory and create temp dir
@@ -112,12 +101,6 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 	// Set project root environment variable
 	os.Setenv("WINDSOR_PROJECT_ROOT", tmpDir)
 
-	// Process options with defaults
-	options := &SetupOptions{}
-	if len(opts) > 0 && opts[0] != nil {
-		options = opts[0]
-	}
-
 	// Create shell
 	mockShell := shell.NewMockShell()
 	// Mock GetProjectRoot to return a temporary directory
@@ -126,12 +109,7 @@ func setupMocks(t *testing.T, opts ...*SetupOptions) *Mocks {
 	}
 
 	// Create config handler
-	var configHandler config.ConfigHandler
-	if options.ConfigHandler == nil {
-		configHandler = config.NewConfigHandler(mockShell)
-	} else {
-		configHandler = options.ConfigHandler
-	}
+	configHandler := config.NewConfigHandler(mockShell)
 
 	// Create mock service
 	mockService := services.NewMockService()
@@ -165,28 +143,30 @@ contexts:
 		t.Fatalf("Failed to load default config string: %v", err)
 	}
 
-	// Load test-specific config if provided
-	if options.ConfigStr != "" {
-		if err := configHandler.LoadConfigString(options.ConfigStr); err != nil {
-			t.Fatalf("Failed to load config string: %v", err)
-		}
-	}
-
 	// Register cleanup to restore original state
 	t.Cleanup(func() {
 		os.Unsetenv("WINDSOR_PROJECT_ROOT")
+		os.Unsetenv("COMPOSE_FILE")
+		os.Unsetenv("WINDSOR_CONTEXT")
 		if err := os.Chdir(origDir); err != nil {
 			t.Logf("Warning: Failed to change back to original directory: %v", err)
 		}
 	})
 
-	return &Mocks{
+	mocks := &VirtTestMocks{
 		Runtime:       rt,
 		ConfigHandler: configHandler,
 		Shell:         mockShell,
 		Service:       mockService,
-		Shims:         setupShims(t),
+		Shims:         setupDefaultShims(),
 	}
+
+	// Apply any overrides
+	for _, opt := range opts {
+		opt(mocks)
+	}
+
+	return mocks
 }
 
 // =============================================================================
@@ -194,9 +174,9 @@ contexts:
 // =============================================================================
 
 func TestVirt_Initialize(t *testing.T) {
-	setup := func(t *testing.T) (*Mocks, *BaseVirt) {
+	setup := func(t *testing.T) (*VirtTestMocks, *BaseVirt) {
 		t.Helper()
-		mocks := setupMocks(t)
+		mocks := setupVirtMocks(t)
 		virt := NewBaseVirt(mocks.Runtime)
 		virt.shims = mocks.Shims
 		return mocks, virt