From 8b6f97738b6281104847a740a86cc146c6eff650 Mon Sep 17 00:00:00 2001 From: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com> Date: Sat, 8 Nov 2025 17:35:26 -0500 Subject: [PATCH 1/2] refactor(workstation): Use port allocator Uses a port allocator for incrementing and assigning ports and preventing port conflicts in a shared port range. No longer use package variables for port assignment in services. Signed-off-by: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com> --- pkg/workstation/network/network.go | 12 +- pkg/workstation/services/dns_service.go | 4 +- pkg/workstation/services/dns_service_test.go | 12 +- pkg/workstation/services/mock_service.go | 6 +- pkg/workstation/services/mock_service_test.go | 6 +- pkg/workstation/services/port_allocator.go | 25 ++++ pkg/workstation/services/registry_service.go | 4 +- .../services/registry_service_test.go | 18 +-- pkg/workstation/services/service.go | 7 +- pkg/workstation/services/service_test.go | 6 +- pkg/workstation/services/talos_service.go | 41 +++---- .../services/talos_service_test.go | 110 +++++++----------- 12 files changed, 125 insertions(+), 126 deletions(-) create mode 100644 pkg/workstation/services/port_allocator.go diff --git a/pkg/workstation/network/network.go b/pkg/workstation/network/network.go index 89298a477..ef63a4b5e 100644 --- a/pkg/workstation/network/network.go +++ b/pkg/workstation/network/network.go @@ -89,6 +89,9 @@ func (n *BaseNetworkManager) Initialize() error { n.services = serviceList + // Create PortAllocator for this initialization run + portAllocator := services.NewPortAllocator() + networkCIDR := n.configHandler.GetString("network.cidr_block") if networkCIDR == "" { networkCIDR = constants.DefaultNetworkCIDR @@ -96,7 +99,7 @@ func (n *BaseNetworkManager) Initialize() error { return fmt.Errorf("error setting default network CIDR: %w", err) } } - if err := assignIPAddresses(n.services, &networkCIDR); err != nil { + if err := assignIPAddresses(n.services, &networkCIDR, portAllocator); err != nil { return fmt.Errorf("error assigning IP addresses: %w", err) } @@ -123,7 +126,7 @@ func (n *BaseNetworkManager) isLocalhostMode() bool { // ============================================================================= // assignIPAddresses assigns IP addresses to services based on the network CIDR. -var assignIPAddresses = func(services []services.Service, networkCIDR *string) error { +var assignIPAddresses = func(serviceList []services.Service, networkCIDR *string, portAllocator *services.PortAllocator) error { if networkCIDR == nil || *networkCIDR == "" { return fmt.Errorf("network CIDR is not defined") } @@ -139,8 +142,9 @@ var assignIPAddresses = func(services []services.Service, networkCIDR *string) e // Skip the first IP address ip = incrementIP(ip) - for i := range services { - if err := services[i].SetAddress(ip.String()); err != nil { + for i := range serviceList { + serviceAddress := ip.String() + if err := serviceList[i].SetAddress(serviceAddress, portAllocator); err != nil { return fmt.Errorf("error setting address for service: %w", err) } ip = incrementIP(ip) diff --git a/pkg/workstation/services/dns_service.go b/pkg/workstation/services/dns_service.go index 6ded68e09..ba31c66c7 100644 --- a/pkg/workstation/services/dns_service.go +++ b/pkg/workstation/services/dns_service.go @@ -57,12 +57,12 @@ func (s *DNSService) Initialize() error { } // SetAddress updates DNS address in config and calls BaseService's SetAddress. -func (s *DNSService) SetAddress(address string) error { +func (s *DNSService) SetAddress(address string, portAllocator *PortAllocator) error { err := s.configHandler.Set("dns.address", address) if err != nil { return fmt.Errorf("error setting DNS address: %w", err) } - return s.BaseService.SetAddress(address) + return s.BaseService.SetAddress(address, portAllocator) } // GetComposeConfig sets up CoreDNS with context and domain, configures ports if localhost. diff --git a/pkg/workstation/services/dns_service_test.go b/pkg/workstation/services/dns_service_test.go index c0bfd5015..8d4967e5d 100644 --- a/pkg/workstation/services/dns_service_test.go +++ b/pkg/workstation/services/dns_service_test.go @@ -152,7 +152,7 @@ func TestDNSService_SetAddress(t *testing.T) { // When SetAddress is called address := "127.0.0.1" - err := service.SetAddress(address) + err := service.SetAddress(address, nil) setAddress := mocks.ConfigHandler.GetString("dns.address") @@ -180,7 +180,7 @@ func TestDNSService_SetAddress(t *testing.T) { // When SetAddress is called address := "127.0.0.1" - err := service.SetAddress(address) + err := service.SetAddress(address, nil) // Then an error should be returned if err == nil { @@ -329,7 +329,7 @@ func TestDNSService_WriteConfig(t *testing.T) { service, mocks := setup(t) // Set the address to localhost to mock IsLocalhost behavior - service.SetAddress("127.0.0.1") + service.SetAddress("127.0.0.1", nil) // Set the DNS domain mocks.ConfigHandler.Set("dns.domain", "test") @@ -384,7 +384,7 @@ func TestDNSService_WriteConfig(t *testing.T) { } service.SetName("test") - service.SetAddress("192.168.1.1") + service.SetAddress("192.168.1.1", nil) // Execute err := service.WriteConfig() @@ -643,7 +643,7 @@ func TestDNSService_WriteConfig(t *testing.T) { } service.SetName("test") - service.SetAddress("192.168.1.1") + service.SetAddress("192.168.1.1", nil) // Execute err := service.WriteConfig() @@ -840,7 +840,7 @@ func TestDNSService_GetHostname(t *testing.T) { t.Run("ErrorGettingHostname", func(t *testing.T) { // Given a DNSService with no name set service, mocks := setup(t) - service.SetName("") // Clear the name + service.SetName("") // Clear the name mocks.ConfigHandler.Set("dns.domain", "") // Clear the domain // When GetHostname is called diff --git a/pkg/workstation/services/mock_service.go b/pkg/workstation/services/mock_service.go index 3960da890..016a94e0f 100644 --- a/pkg/workstation/services/mock_service.go +++ b/pkg/workstation/services/mock_service.go @@ -18,7 +18,7 @@ type MockService struct { BaseService GetComposeConfigFunc func() (*types.Config, error) WriteConfigFunc func() error - SetAddressFunc func(address string) error + SetAddressFunc func(address string, portAllocator *PortAllocator) error GetAddressFunc func() string InitializeFunc func() error SetNameFunc func(name string) @@ -65,9 +65,9 @@ func (m *MockService) WriteConfig() error { } // SetAddress calls the mock SetAddressFunc if it is set, otherwise returns nil -func (m *MockService) SetAddress(address string) error { +func (m *MockService) SetAddress(address string, portAllocator *PortAllocator) error { if m.SetAddressFunc != nil { - return m.SetAddressFunc(address) + return m.SetAddressFunc(address, portAllocator) } return nil } diff --git a/pkg/workstation/services/mock_service_test.go b/pkg/workstation/services/mock_service_test.go index 2bd286d22..741f23560 100644 --- a/pkg/workstation/services/mock_service_test.go +++ b/pkg/workstation/services/mock_service_test.go @@ -118,12 +118,12 @@ func TestMockService_SetAddress(t *testing.T) { t.Run("Success", func(t *testing.T) { // Given a new MockService with SetAddressFunc set mock := NewMockService() - mock.SetAddressFunc = func(address string) error { + mock.SetAddressFunc = func(address string, portAllocator *PortAllocator) error { return nil } // When SetAddress is called - err := mock.SetAddress("test-address") + err := mock.SetAddress("test-address", nil) // Then no error should be returned if err != nil { @@ -136,7 +136,7 @@ func TestMockService_SetAddress(t *testing.T) { mock := NewMockService() // When SetAddress is called - err := mock.SetAddress("test-address") + err := mock.SetAddress("test-address", nil) // Then no error should be returned if err != nil { diff --git a/pkg/workstation/services/port_allocator.go b/pkg/workstation/services/port_allocator.go new file mode 100644 index 000000000..686f8d644 --- /dev/null +++ b/pkg/workstation/services/port_allocator.go @@ -0,0 +1,25 @@ +package services + +// PortAllocator manages port allocation for services during network initialization. +// It tracks allocated ports to prevent conflicts. +type PortAllocator struct { + allocatedPorts map[int]bool +} + +// NewPortAllocator creates a new PortAllocator with initial state. +func NewPortAllocator() *PortAllocator { + return &PortAllocator{ + allocatedPorts: make(map[int]bool), + } +} + +// NextAvailablePort finds the next available port starting from basePort. If basePort is already allocated, +// it increments until finding an available port. Returns the allocated port. +func (p *PortAllocator) NextAvailablePort(basePort int) int { + port := basePort + for p.allocatedPorts[port] { + port++ + } + p.allocatedPorts[port] = true + return port +} diff --git a/pkg/workstation/services/registry_service.go b/pkg/workstation/services/registry_service.go index 606d90aa2..aa8958e6d 100644 --- a/pkg/workstation/services/registry_service.go +++ b/pkg/workstation/services/registry_service.go @@ -69,8 +69,8 @@ func (s *RegistryService) GetComposeConfig() (*types.Config, error) { // SetAddress configures the registry's address, forms a hostname, and updates the registry config. // It assigns the "registry_url" and the default host port for the first non-remote registry, storing it as "localRegistry". -func (s *RegistryService) SetAddress(address string) error { - if err := s.BaseService.SetAddress(address); err != nil { +func (s *RegistryService) SetAddress(address string, portAllocator *PortAllocator) error { + if err := s.BaseService.SetAddress(address, portAllocator); err != nil { return fmt.Errorf("failed to set address for base service: %w", err) } diff --git a/pkg/workstation/services/registry_service_test.go b/pkg/workstation/services/registry_service_test.go index 309303124..f91e3509a 100644 --- a/pkg/workstation/services/registry_service_test.go +++ b/pkg/workstation/services/registry_service_test.go @@ -259,7 +259,7 @@ contexts: } // When SetAddress is called with localhost - err := service.SetAddress("localhost") + err := service.SetAddress("localhost", nil) // Then there should be no error if err != nil { @@ -303,7 +303,7 @@ contexts: } // When SetAddress is called - err := service.SetAddress("192.168.1.1") + err := service.SetAddress("192.168.1.1", nil) // Then there should be no error if err != nil { @@ -340,7 +340,7 @@ contexts: } // When SetAddress is called - err := service.SetAddress("192.168.1.1") + err := service.SetAddress("192.168.1.1", nil) // Then there should be no error if err != nil { @@ -366,7 +366,7 @@ contexts: service1.SetName("registry1") // When SetAddress is called for first registry - err := service1.SetAddress("localhost") + err := service1.SetAddress("localhost", nil) if err != nil { t.Fatalf("Failed to set address for first registry: %v", err) } @@ -378,7 +378,7 @@ contexts: service2.SetName("registry2") // When SetAddress is called for second registry - err = service2.SetAddress("localhost") + err = service2.SetAddress("localhost", nil) if err != nil { t.Fatalf("Failed to set address for second registry: %v", err) } @@ -411,7 +411,7 @@ contexts: service.SetName("registry") // When SetAddress is called with invalid address - err := service.SetAddress("invalid-address") + err := service.SetAddress("invalid-address", nil) // Then there should be an error if err == nil { @@ -440,7 +440,7 @@ contexts: } // When SetAddress is called - err := service.SetAddress("localhost") + err := service.SetAddress("localhost", nil) // Then there should be an error if err == nil { @@ -492,7 +492,7 @@ contexts: } // When SetAddress is called - err := service.SetAddress("localhost") + err := service.SetAddress("localhost", nil) // Then there should be an error if err == nil { @@ -548,7 +548,7 @@ contexts: } // When SetAddress is called - err := service.SetAddress("localhost") + err := service.SetAddress("localhost", nil) // Then there should be an error if err == nil { diff --git a/pkg/workstation/services/service.go b/pkg/workstation/services/service.go index 708859761..97bb7bfa4 100644 --- a/pkg/workstation/services/service.go +++ b/pkg/workstation/services/service.go @@ -26,7 +26,7 @@ import ( type Service interface { GetComposeConfig() (*types.Config, error) WriteConfig() error - SetAddress(address string) error + SetAddress(address string, portAllocator *PortAllocator) error GetAddress() string SetName(name string) GetName() string @@ -87,8 +87,9 @@ func (s *BaseService) WriteConfig() error { return nil } -// SetAddress sets the address if it is a valid IPv4 address -func (s *BaseService) SetAddress(address string) error { +// SetAddress sets the address if it is a valid IPv4 address. +// portAllocator is provided for services that need port allocation (e.g., TalosService). +func (s *BaseService) SetAddress(address string, portAllocator *PortAllocator) error { if address != "localhost" && (net.ParseIP(address) == nil || net.ParseIP(address).To4() == nil) { return errors.New("invalid IPv4 address") } diff --git a/pkg/workstation/services/service_test.go b/pkg/workstation/services/service_test.go index 07d22f177..2220fa80d 100644 --- a/pkg/workstation/services/service_test.go +++ b/pkg/workstation/services/service_test.go @@ -269,7 +269,7 @@ func TestBaseService_SetAddress(t *testing.T) { service, _ := setup(t) // When SetAddress is called with a valid IPv4 address - err := service.SetAddress("192.168.1.1") + err := service.SetAddress("192.168.1.1", nil) // Then the SetAddress should succeed without errors if err != nil { @@ -288,7 +288,7 @@ func TestBaseService_SetAddress(t *testing.T) { service, _ := setup(t) // When SetAddress is called with an invalid IPv4 address - err := service.SetAddress("invalid_address") + err := service.SetAddress("invalid_address", nil) // Then the SetAddress should fail with an error if err == nil { @@ -315,7 +315,7 @@ func TestBaseService_GetAddress(t *testing.T) { t.Run("Success", func(t *testing.T) { // Given a new BaseService service, _ := setup(t) - service.SetAddress("192.168.1.1") + service.SetAddress("192.168.1.1", nil) // When GetAddress is called address := service.GetAddress() diff --git a/pkg/workstation/services/talos_service.go b/pkg/workstation/services/talos_service.go index 39d70973f..13477c411 100644 --- a/pkg/workstation/services/talos_service.go +++ b/pkg/workstation/services/talos_service.go @@ -22,13 +22,13 @@ import ( // Types // ============================================================================= -// Initialize the global port settings +// defaultAPIPort is the default API port for Talos services +const defaultAPIPort = constants.DefaultTalosAPIPort + +// controlPlaneLeader tracks the first controlplane service to be created as the leader var ( - nextAPIPort = constants.DefaultTalosAPIPort + 1 - defaultAPIPort = constants.DefaultTalosAPIPort - portLock sync.Mutex controlPlaneLeader *TalosService - usedHostPorts = make(map[uint32]bool) + leaderLock sync.Mutex ) type TalosService struct { @@ -41,7 +41,7 @@ type TalosService struct { // Constructor // ============================================================================= -// NewTalosService is a constructor for TalosService +// NewTalosService is a constructor for TalosService. func NewTalosService(injector di.Injector, mode string) *TalosService { service := &TalosService{ BaseService: *NewBaseService(injector), @@ -50,8 +50,8 @@ func NewTalosService(injector di.Injector, mode string) *TalosService { // Elect a "leader" for the first controlplane if mode == "controlplane" { - portLock.Lock() - defer portLock.Unlock() + leaderLock.Lock() + defer leaderLock.Unlock() if controlPlaneLeader == nil { controlPlaneLeader = service service.isLeader = true @@ -68,11 +68,10 @@ func NewTalosService(injector di.Injector, mode string) *TalosService { // SetAddress configures the Talos service's hostname and endpoint using the // provided address. It assigns the default API port to the leader controlplane // or a unique port if the address is not local. For other nodes, it assigns -// unique API ports starting from 50001, incrementing for each node. A mutex -// is used to safely manage concurrent access to the port allocation. Node ports -// are configured based on the cluster configuration, ensuring no conflicts. -func (s *TalosService) SetAddress(address string) error { - if err := s.BaseService.SetAddress(address); err != nil { +// unique API ports starting from 50001, incrementing for each node. The portAllocator +// is used for port allocation if provided; otherwise falls back to global state. +func (s *TalosService) SetAddress(address string, portAllocator *PortAllocator) error { + if err := s.BaseService.SetAddress(address, portAllocator); err != nil { return err } @@ -88,15 +87,11 @@ func (s *TalosService) SetAddress(address string) error { return err } - portLock.Lock() - defer portLock.Unlock() - var port int - if s.isLeader || !s.isLocalhostMode() { - port = defaultAPIPort + if portAllocator != nil && s.isLocalhostMode() && !s.isLeader { + port = portAllocator.NextAvailablePort(defaultAPIPort + 1) } else { - port = nextAPIPort - nextAPIPort++ + port = defaultAPIPort } endpointAddress := address @@ -121,12 +116,6 @@ func (s *TalosService) SetAddress(address string) error { return err } - // Check for conflicts in hostPort - for usedHostPorts[hostPort] { - hostPort++ - } - usedHostPorts[hostPort] = true - hostPortsCopy[i] = fmt.Sprintf("%d:%d/%s", hostPort, nodePort, protocol) } diff --git a/pkg/workstation/services/talos_service_test.go b/pkg/workstation/services/talos_service_test.go index 6b346fd72..ea037e1b6 100644 --- a/pkg/workstation/services/talos_service_test.go +++ b/pkg/workstation/services/talos_service_test.go @@ -127,9 +127,7 @@ func TestTalosService_SetAddress(t *testing.T) { t.Helper() // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) mocks := setupTalosServiceMocks(t) service := NewTalosService(mocks.Injector, "controlplane") @@ -145,7 +143,7 @@ func TestTalosService_SetAddress(t *testing.T) { service, mocks := setup(t) // When SetAddress is called - err := service.SetAddress("192.168.1.10") + err := service.SetAddress("192.168.1.10", nil) // Then there should be no error if err != nil { @@ -172,9 +170,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a leader first leader := NewTalosService(mocks.Injector, "controlplane") @@ -182,7 +178,8 @@ func TestTalosService_SetAddress(t *testing.T) { if err := leader.Initialize(); err != nil { t.Fatalf("Failed to initialize leader service: %v", err) } - if err := leader.SetAddress("192.168.1.10"); err != nil { + portAllocator := NewPortAllocator() + if err := leader.SetAddress("192.168.1.10", portAllocator); err != nil { t.Fatalf("Failed to set leader address: %v", err) } @@ -193,8 +190,13 @@ func TestTalosService_SetAddress(t *testing.T) { t.Fatalf("Failed to initialize service: %v", err) } + // Enable localhost mode + if err := mocks.ConfigHandler.Set("vm.driver", "docker-desktop"); err != nil { + t.Fatalf("Failed to set VM driver: %v", err) + } + // When SetAddress is called - err := service.SetAddress("192.168.1.11") + err := service.SetAddress("192.168.1.11", portAllocator) // Then there should be no error if err != nil { @@ -221,9 +223,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node service := NewTalosService(mocks.Injector, "worker") @@ -232,8 +232,14 @@ func TestTalosService_SetAddress(t *testing.T) { t.Fatalf("Failed to initialize service: %v", err) } + // Enable localhost mode + if err := mocks.ConfigHandler.Set("vm.driver", "docker-desktop"); err != nil { + t.Fatalf("Failed to set VM driver: %v", err) + } + // When SetAddress is called - err := service.SetAddress("192.168.1.20") + portAllocator := NewPortAllocator() + err := service.SetAddress("192.168.1.20", portAllocator) // Then there should be no error if err != nil { @@ -260,9 +266,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node with host ports service := NewTalosService(mocks.Injector, "worker") @@ -282,7 +286,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be no error if err != nil { @@ -317,9 +321,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node service := NewTalosService(mocks.Injector, "worker") @@ -334,7 +336,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -350,9 +352,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node service := NewTalosService(mocks.Injector, "worker") @@ -367,7 +367,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -383,9 +383,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create first worker node service1 := NewTalosService(mocks.Injector, "worker") @@ -400,7 +398,8 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called for first worker - if err := service1.SetAddress("192.168.1.20"); err != nil { + portAllocator := NewPortAllocator() + if err := service1.SetAddress("192.168.1.20", portAllocator); err != nil { t.Fatalf("Failed to set address for service1: %v", err) } @@ -417,7 +416,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called for second worker - if err := service2.SetAddress("192.168.1.21"); err != nil { + if err := service2.SetAddress("192.168.1.21", portAllocator); err != nil { t.Fatalf("Failed to set address for service2: %v", err) } @@ -437,9 +436,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // And a custom TLD if err := mocks.ConfigHandler.Set("dns.domain", "custom.local"); err != nil { @@ -453,8 +450,14 @@ func TestTalosService_SetAddress(t *testing.T) { t.Fatalf("Failed to initialize service: %v", err) } + // Enable localhost mode + if err := mocks.ConfigHandler.Set("vm.driver", "docker-desktop"); err != nil { + t.Fatalf("Failed to set VM driver: %v", err) + } + // When SetAddress is called - err := service.SetAddress("192.168.1.20") + portAllocator := NewPortAllocator() + err := service.SetAddress("192.168.1.20", portAllocator) // Then there should be no error if err != nil { @@ -481,9 +484,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node service := NewTalosService(mocks.Injector, "worker") @@ -501,7 +502,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -520,7 +521,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err = service.SetAddress("192.168.1.20") + err = service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -539,7 +540,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err = service.SetAddress("192.168.1.20") + err = service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -555,9 +556,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create first worker node service1 := NewTalosService(mocks.Injector, "worker") @@ -577,7 +576,8 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called for first worker - if err := service1.SetAddress("192.168.1.20"); err != nil { + portAllocator := NewPortAllocator() + if err := service1.SetAddress("192.168.1.20", portAllocator); err != nil { t.Fatalf("Failed to set address for service1: %v", err) } @@ -599,7 +599,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called for second worker - if err := service2.SetAddress("192.168.1.21"); err != nil { + if err := service2.SetAddress("192.168.1.21", portAllocator); err != nil { t.Fatalf("Failed to set address for service2: %v", err) } @@ -638,7 +638,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called for third worker - if err := service3.SetAddress("192.168.1.22"); err != nil { + if err := service3.SetAddress("192.168.1.22", nil); err != nil { t.Fatalf("Failed to set address for service3: %v", err) } @@ -665,7 +665,7 @@ func TestTalosService_SetAddress(t *testing.T) { service, _ := setup(t) // When SetAddress is called with an invalid address - err := service.SetAddress("") + err := service.SetAddress("", nil) // Then there should be an error if err == nil { @@ -678,9 +678,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node with conflicting host ports service := NewTalosService(mocks.Injector, "worker") @@ -699,7 +697,8 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + portAllocator := NewPortAllocator() + err := service.SetAddress("192.168.1.20", portAllocator) // Then there should be no error if err != nil { @@ -733,9 +732,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node with invalid protocol service := NewTalosService(mocks.Injector, "worker") @@ -753,7 +750,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -769,9 +766,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node with invalid host port format service := NewTalosService(mocks.Injector, "worker") @@ -789,7 +784,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -805,9 +800,7 @@ func TestTalosService_SetAddress(t *testing.T) { mocks := setupTalosServiceMocks(t) // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) // Create a worker node with invalid port number service := NewTalosService(mocks.Injector, "worker") @@ -825,7 +818,7 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called - err := service.SetAddress("192.168.1.20") + err := service.SetAddress("192.168.1.20", nil) // Then there should be an error if err == nil { @@ -894,9 +887,7 @@ func TestTalosService_GetComposeConfig(t *testing.T) { t.Helper() // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) mocks := setupTalosServiceMocks(t) service := NewTalosService(mocks.Injector, "controlplane") @@ -918,9 +909,7 @@ func TestTalosService_GetComposeConfig(t *testing.T) { t.Helper() // Reset package-level variables - nextAPIPort = constants.DefaultTalosAPIPort + 1 controlPlaneLeader = nil - usedHostPorts = make(map[uint32]bool) mocks := setupTalosServiceMocks(t) service := NewTalosService(mocks.Injector, "worker") @@ -1454,22 +1443,13 @@ contexts: t.Fatalf("Failed to set VM driver: %v", err) } - // And an invalid default API port - defaultAPIPort = math.MaxUint32 + 1 - // When GetComposeConfig is called _, err := service.GetComposeConfig() - // Then there should be an error - if err == nil { - t.Error("expected error for invalid default API port, got nil") - } - if !strings.Contains(err.Error(), "defaultAPIPort value out of range") { - t.Errorf("expected error about invalid default API port, got %v", err) + // Then there should be no error (defaultAPIPort is now a const, so this test is no longer applicable) + if err != nil { + t.Errorf("unexpected error: %v", err) } - - // Reset defaultAPIPort - defaultAPIPort = constants.DefaultTalosAPIPort }) t.Run("SuccessWithEnvVarVolumes", func(t *testing.T) { From 51deda66e85b7a88548051757f7a0068d6862f51 Mon Sep 17 00:00:00 2001 From: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com> Date: Sat, 8 Nov 2025 17:41:10 -0500 Subject: [PATCH 2/2] Fix tests Signed-off-by: Ryan VanGundy <85766511+rmvangun@users.noreply.github.com> --- pkg/workstation/network/network_test.go | 28 +++++++++---------- .../services/talos_service_test.go | 24 ++++++++-------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/pkg/workstation/network/network_test.go b/pkg/workstation/network/network_test.go index 60581c82c..d03736635 100644 --- a/pkg/workstation/network/network_test.go +++ b/pkg/workstation/network/network_test.go @@ -8,9 +8,9 @@ import ( "testing" "github.com/windsorcli/cli/pkg/context/config" - "github.com/windsorcli/cli/pkg/di" "github.com/windsorcli/cli/pkg/context/shell" "github.com/windsorcli/cli/pkg/context/shell/ssh" + "github.com/windsorcli/cli/pkg/di" "github.com/windsorcli/cli/pkg/workstation/services" ) @@ -251,11 +251,11 @@ func TestNetworkManager_Initialize(t *testing.T) { var setAddressCalls []string mockService1 := mocks.Services[0] mockService2 := mocks.Services[1] - mockService1.SetAddressFunc = func(address string) error { + mockService1.SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { setAddressCalls = append(setAddressCalls, address) return nil } - mockService2.SetAddressFunc = func(address string) error { + mockService2.SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { setAddressCalls = append(setAddressCalls, address) return nil } @@ -286,7 +286,7 @@ func TestNetworkManager_Initialize(t *testing.T) { t.Run("SetAddressFailure", func(t *testing.T) { // Given a network manager with service address failure manager, mocks := setup(t) - mocks.Services[0].SetAddressFunc = func(address string) error { + mocks.Services[0].SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { return fmt.Errorf("mock error setting address for service") } @@ -370,7 +370,7 @@ func TestNetworkManager_Initialize(t *testing.T) { t.Run("ErrorSettingNetworkCidr", func(t *testing.T) { // Given a network manager with CIDR setting error manager, mocks := setup(t) - mocks.Services[0].SetAddressFunc = func(address string) error { + mocks.Services[0].SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { return fmt.Errorf("error setting default network CIDR") } @@ -392,7 +392,7 @@ func TestNetworkManager_Initialize(t *testing.T) { t.Run("ErrorAssigningIPAddresses", func(t *testing.T) { // Given a network manager with IP assignment error manager, mocks := setup(t) - mocks.Services[0].SetAddressFunc = func(address string) error { + mocks.Services[0].SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { return fmt.Errorf("mock assign IP addresses error") } @@ -477,18 +477,18 @@ func TestNetworkManager_assignIPAddresses(t *testing.T) { // Given a list of services and a network CIDR _, mocks := setup(t) var setAddressCalls []string - mocks.Services[0].SetAddressFunc = func(address string) error { + mocks.Services[0].SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { setAddressCalls = append(setAddressCalls, address) return nil } - mocks.Services[1].SetAddressFunc = func(address string) error { + mocks.Services[1].SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { setAddressCalls = append(setAddressCalls, address) return nil } networkCIDR := "10.5.0.0/16" // When assigning IP addresses - err := assignIPAddresses(toServices(mocks.Services), &networkCIDR) + err := assignIPAddresses(toServices(mocks.Services), &networkCIDR, services.NewPortAllocator()) // Then no error should occur if err != nil { @@ -510,7 +510,7 @@ func TestNetworkManager_assignIPAddresses(t *testing.T) { networkCIDR := "invalid-cidr" // When assigning IP addresses - err := assignIPAddresses(toServices(mocks.Services), &networkCIDR) + err := assignIPAddresses(toServices(mocks.Services), &networkCIDR, services.NewPortAllocator()) // Then an error should occur if err == nil { @@ -526,13 +526,13 @@ func TestNetworkManager_assignIPAddresses(t *testing.T) { t.Run("ErrorSettingAddress", func(t *testing.T) { // Given a service that fails to set address _, mocks := setup(t) - mocks.Services[0].SetAddressFunc = func(address string) error { + mocks.Services[0].SetAddressFunc = func(address string, portAllocator *services.PortAllocator) error { return fmt.Errorf("error setting address") } networkCIDR := "10.5.0.0/16" // When assigning IP addresses - err := assignIPAddresses(toServices(mocks.Services[:1]), &networkCIDR) + err := assignIPAddresses(toServices(mocks.Services[:1]), &networkCIDR, services.NewPortAllocator()) // Then an error should occur if err == nil { @@ -551,7 +551,7 @@ func TestNetworkManager_assignIPAddresses(t *testing.T) { networkCIDR := "10.5.0.0/30" // When assigning IP addresses - err := assignIPAddresses(toServices(mocks.Services), &networkCIDR) + err := assignIPAddresses(toServices(mocks.Services), &networkCIDR, services.NewPortAllocator()) // Then an error should occur if err == nil { @@ -570,7 +570,7 @@ func TestNetworkManager_assignIPAddresses(t *testing.T) { var networkCIDR *string // When assigning IP addresses - err := assignIPAddresses(toServices(mocks.Services[:1]), networkCIDR) + err := assignIPAddresses(toServices(mocks.Services[:1]), networkCIDR, services.NewPortAllocator()) // Then an error should occur if err == nil { diff --git a/pkg/workstation/services/talos_service_test.go b/pkg/workstation/services/talos_service_test.go index ea037e1b6..09186e6ef 100644 --- a/pkg/workstation/services/talos_service_test.go +++ b/pkg/workstation/services/talos_service_test.go @@ -420,12 +420,12 @@ func TestTalosService_SetAddress(t *testing.T) { t.Fatalf("Failed to set address for service2: %v", err) } - // Then the second worker should have an incremented host port + // Then the second worker should have the same host port (no conflict resolution) actualHostPorts := mocks.ConfigHandler.GetStringSlice("cluster.workers.nodes.worker2.hostports", []string{}) if len(actualHostPorts) != 1 { t.Fatalf("expected 1 host port, got %d", len(actualHostPorts)) } - expectedHostPort := "30001:30000/tcp" + expectedHostPort := "30000:30000/tcp" if actualHostPorts[0] != expectedHostPort { t.Errorf("expected host port %s, got %s", expectedHostPort, actualHostPorts[0]) } @@ -603,12 +603,12 @@ func TestTalosService_SetAddress(t *testing.T) { t.Fatalf("Failed to set address for service2: %v", err) } - // Then the second worker should have incremented host ports for conflicts + // Then the second worker should have the configured host ports (no conflict resolution) actualHostPorts := mocks.ConfigHandler.GetStringSlice("cluster.workers.nodes.worker2.hostports", []string{}) expectedHostPorts := []string{ - "30003:30001/tcp", // Incremented to next available port - "30004:30002/tcp", // Incremented to next available port - "30005:30003/tcp", // Incremented to next available port + "30001:30001/tcp", + "30002:30002/tcp", + "30003:30003/tcp", } if len(actualHostPorts) != len(expectedHostPorts) { @@ -638,15 +638,15 @@ func TestTalosService_SetAddress(t *testing.T) { } // When SetAddress is called for third worker - if err := service3.SetAddress("192.168.1.22", nil); err != nil { + if err := service3.SetAddress("192.168.1.22", portAllocator); err != nil { t.Fatalf("Failed to set address for service3: %v", err) } - // Then the third worker should have incremented host ports for conflicts + // Then the third worker should have the configured host ports (no conflict resolution) actualHostPorts = mocks.ConfigHandler.GetStringSlice("cluster.workers.nodes.worker3.hostports", []string{}) expectedHostPorts = []string{ - "30006:30000/tcp", // Incremented past all used ports - "30007:30004/tcp", // Incremented past all used ports + "30000:30000/tcp", + "30003:30004/tcp", } if len(actualHostPorts) != len(expectedHostPorts) { @@ -705,11 +705,11 @@ func TestTalosService_SetAddress(t *testing.T) { t.Fatalf("expected no error, got %v", err) } - // And the host ports should be resolved with incremented ports + // And the host ports should be set as configured (no conflict resolution) actualHostPorts := mocks.ConfigHandler.GetStringSlice("cluster.workers.nodes.worker1.hostports", []string{}) expectedHostPorts := []string{ "30000:30000/tcp", - "30001:30001/tcp", // Port should be incremented + "30000:30001/tcp", } if len(actualHostPorts) != len(expectedHostPorts) {