103 lines
2.1 KiB
Go
103 lines
2.1 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"k8s.io/apimachinery/pkg/watch"
|
|
)
|
|
|
|
// MockDiscovery is an in-memory Discovery implementation for tests.
|
|
// It allows deterministic injection of host additions/removals without
|
|
// depending on Kubernetes API machinery.
|
|
type MockDiscovery struct {
|
|
mu sync.RWMutex
|
|
hosts []string
|
|
events chan HostChange
|
|
closed bool
|
|
started bool
|
|
}
|
|
|
|
// NewMockDiscovery creates a mock discovery with an initial host list.
|
|
func NewMockDiscovery(initial []string) *MockDiscovery {
|
|
cp := make([]string, len(initial))
|
|
copy(cp, initial)
|
|
return &MockDiscovery{
|
|
hosts: cp,
|
|
events: make(chan HostChange, 32),
|
|
}
|
|
}
|
|
|
|
// Discover returns the current host snapshot.
|
|
func (m *MockDiscovery) Discover() ([]string, error) {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
cp := make([]string, len(m.hosts))
|
|
copy(cp, m.hosts)
|
|
return cp, nil
|
|
}
|
|
|
|
// Watch returns a channel that will receive HostChange events.
|
|
// The channel is buffered; AddHost/RemoveHost push events non-blockingly.
|
|
func (m *MockDiscovery) Watch() (<-chan HostChange, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return nil, context.Canceled
|
|
}
|
|
m.started = true
|
|
return m.events, nil
|
|
}
|
|
|
|
// AddHost inserts a new host (if absent) and emits an Added event.
|
|
func (m *MockDiscovery) AddHost(host string) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return
|
|
}
|
|
for _, h := range m.hosts {
|
|
if h == host {
|
|
return
|
|
}
|
|
}
|
|
m.hosts = append(m.hosts, host)
|
|
if m.started {
|
|
m.events <- HostChange{Host: host, Type: watch.Added}
|
|
}
|
|
}
|
|
|
|
// RemoveHost removes a host (if present) and emits a Deleted event.
|
|
func (m *MockDiscovery) RemoveHost(host string) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return
|
|
}
|
|
idx := -1
|
|
for i, h := range m.hosts {
|
|
if h == host {
|
|
idx = i
|
|
break
|
|
}
|
|
}
|
|
if idx == -1 {
|
|
return
|
|
}
|
|
m.hosts = append(m.hosts[:idx], m.hosts[idx+1:]...)
|
|
if m.started {
|
|
m.events <- HostChange{Host: host, Type: watch.Deleted}
|
|
}
|
|
}
|
|
|
|
// Close closes the event channel (idempotent).
|
|
func (m *MockDiscovery) Close() {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.closed {
|
|
return
|
|
}
|
|
m.closed = true
|
|
close(m.events)
|
|
}
|