mocks.go

/*
Copyright 2025 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package mocks provides mocks for the interfaces defined in the `contracts` package.
//
// # Testing Philosophy: High-Fidelity Mocks
//
// The components that consume these contracts, particularly the `controller.Processor`, are complex, concurrent
// orchestrators. Testing them reliably requires more than simple stubs. It requires high-fidelity mocks that allow for
// the deterministic simulation of race conditions and specific failure modes.
//
// For this reason, mocks like `MockManagedQueue` are deliberately stateful and thread-safe. They provide a reliable,
// in-memory simulation of the real component's behavior, while also providing function-based overrides
// (e.g., `AddFunc`) that allow tests to inject specific errors or pause execution at critical moments. This strategy is
// essential for creating the robust, non-flaky tests needed to verify the correctness of the system's concurrent logic.
// For a more detailed defense of this strategy, see the comment at the top of `controller/internal/processor_test.go`.
package mocks

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/llm-d/llm-d-router/pkg/epp/flowcontrol/contracts"
	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol"
	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol/mocks"
)

// --- RegistryDataPlane Mocks ---

// MockRegistryDataPlane is a simple "stub-style" mock for testing.
// Its methods are implemented as function fields (e.g., `IDFunc`). A test can inject behavior by setting the desired
// function field in the test setup. If a func is nil, the method will return a zero value.
type MockRegistryDataPlane struct {
	ManagedQueueFunc             func(key flowcontrol.FlowKey) (contracts.ManagedQueue, error)
	FairnessPolicyFunc           func(priority int) (flowcontrol.FairnessPolicy, error)
	PriorityBandAccessorFunc     func(priority int) (flowcontrol.PriorityBandAccessor, error)
	AllOrderedPriorityLevelsFunc func() []int
	StatsFunc                    func() contracts.AggregateStats
	WithConnectionFunc           func(key flowcontrol.FlowKey, fn func(conn contracts.ActiveFlowConnection) error) error
}

func (m *MockRegistryDataPlane) ManagedQueue(key flowcontrol.FlowKey) (contracts.ManagedQueue, error) {
	if m.ManagedQueueFunc != nil {
		return m.ManagedQueueFunc(key)
	}
	return &MockManagedQueue{FlowKeyV: key}, nil
}

func (m *MockRegistryDataPlane) FairnessPolicy(priority int) (flowcontrol.FairnessPolicy, error) {
	if m.FairnessPolicyFunc != nil {
		return m.FairnessPolicyFunc(priority)
	}
	return nil, errors.New("sentinel error for mock data plane")
}

func (m *MockRegistryDataPlane) PriorityBandAccessor(priority int) (flowcontrol.PriorityBandAccessor, error) {
	if m.PriorityBandAccessorFunc != nil {
		return m.PriorityBandAccessorFunc(priority)
	}
	return nil, errors.New("sentinel error for mock data plane")
}

func (m *MockRegistryDataPlane) AllOrderedPriorityLevels() []int {
	if m.AllOrderedPriorityLevelsFunc != nil {
		return m.AllOrderedPriorityLevelsFunc()
	}
	return nil
}

func (m *MockRegistryDataPlane) Stats() contracts.AggregateStats {
	if m.StatsFunc != nil {
		return m.StatsFunc()
	}
	return contracts.AggregateStats{}
}

func (m *MockRegistryDataPlane) WithConnection(key flowcontrol.FlowKey, fn func(conn contracts.ActiveFlowConnection) error) error {
	if m.WithConnectionFunc != nil {
		return m.WithConnectionFunc(key, fn)
	}
	return nil
}

func (m *MockRegistryDataPlane) SubmitDesiredPriorities(_ map[int]struct{}) {}

func (m *MockRegistryDataPlane) PriorityBandUpdateChannel() <-chan map[int]struct{} {
	return nil
}

func (m *MockRegistryDataPlane) FlowGCTimeout() time.Duration {
	return time.Minute
}

func (m *MockRegistryDataPlane) ApplyDesiredPriorities(_ map[int]struct{}) {}

func (m *MockRegistryDataPlane) ExecuteGCCycle() {}

var _ contracts.FlowRegistry = &MockRegistryDataPlane{}

var _ contracts.FlowRegistryDataPlane = &MockRegistryDataPlane{}

// --- Dependency Mocks ---

// MockSaturationDetector is a simple "stub-style" mock for testing.
type MockSaturationDetector struct {
	SaturationFunc func(ctx context.Context, candidatePods []fwkdl.Endpoint) float64
}

func (m *MockSaturationDetector) Saturation(ctx context.Context, candidatePods []fwkdl.Endpoint) float64 {
	if m.SaturationFunc != nil {
		return m.SaturationFunc(ctx, candidatePods)
	}
	return 0.0
}

// MockEndpointCandidates provides a mock implementation of the contracts.EndpointCandidates interface.
// It allows tests to control the exact set of endpoint candidates returned for a given request.
type MockEndpointCandidates struct {
	// LocateFunc allows injecting custom logic.
	LocateFunc func(ctx context.Context, requestMetadata map[string]any) []fwkdl.Endpoint
	// Candidates is a static return value used if LocateFunc is nil.
	Candidates []fwkdl.Endpoint
}

func (m *MockEndpointCandidates) Locate(ctx context.Context, requestMetadata map[string]any) []fwkdl.Endpoint {
	if m.LocateFunc != nil {
		return m.LocateFunc(ctx, requestMetadata)
	}
	// Return copy to be safe
	if m.Candidates == nil {
		return nil
	}

	result := make([]fwkdl.Endpoint, len(m.Candidates))
	copy(result, m.Candidates)
	return result
}

// --- SafeQueue Mock ---

// MockSafeQueue is a simple stub mock for the SafeQueue interface.
// It is used for tests that need to control the exact return values of a queue's methods without simulating the queue's
// internal logic or state.
type MockSafeQueue struct {
	NameV       string
	LenV        int
	ByteSizeV   uint64
	PeekV       flowcontrol.QueueItemAccessor
	AddFunc     func(item flowcontrol.QueueItemAccessor)
	RemoveFunc  func(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error)
	CleanupFunc func(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor
	DrainFunc   func() []flowcontrol.QueueItemAccessor
}

func (m *MockSafeQueue) Name() string     { return m.NameV }
func (m *MockSafeQueue) Len() int         { return m.LenV }
func (m *MockSafeQueue) ByteSize() uint64 { return m.ByteSizeV }

func (m *MockSafeQueue) Peek() flowcontrol.QueueItemAccessor {
	return m.PeekV
}

func (m *MockSafeQueue) Add(item flowcontrol.QueueItemAccessor) {
	if m.AddFunc != nil {
		m.AddFunc(item)
	}
}

func (m *MockSafeQueue) Remove(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error) {
	if m.RemoveFunc != nil {
		return m.RemoveFunc(handle)
	}
	return nil, errors.New("sentinel error for mock queue")
}

func (m *MockSafeQueue) Cleanup(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor {
	if m.CleanupFunc != nil {
		return m.CleanupFunc(predicate)
	}
	return nil
}

func (m *MockSafeQueue) Drain() []flowcontrol.QueueItemAccessor {
	if m.DrainFunc != nil {
		return m.DrainFunc()
	}
	return nil
}

var _ contracts.SafeQueue = &MockSafeQueue{}

// --- ManagedQueue Mock ---

// MockManagedQueue is a high-fidelity, thread-safe mock of the `contracts.ManagedQueue` interface, designed
// specifically for testing the concurrent `controller/internal.Processor`.
//
// This mock is essential for creating deterministic and focused unit tests. It allows for precise control over queue
// behavior and enables the testing of critical edge cases (e.g., empty queues, dispatch failures) in complete
// isolation, which would be difficult and unreliable to achieve with the concrete `registry.managedQueue`
// implementation.
//
// ### Design Philosophy
//
//  1. **Stateful**: The mock maintains an internal map of items to accurately reflect a real queue's state. Its `Len()`
//     and `ByteSize()` methods are derived directly from this state.
//  2. **Deadlock-Safe Overrides**: Test-specific logic (e.g., `AddFunc`) is executed instead of the default
//     implementation. The override function is fully responsible for its own logic and synchronization, as the mock's
//     internal mutex will *not* be held during its execution.
//  3. **Self-Wiring**: The `FlowQueueAccessor()` method returns the mock itself, ensuring the accessor is always
//     correctly connected to the queue's state without manual wiring in tests.
type MockManagedQueue struct {
	// FlowKeyV defines the flow specification for this mock queue. It should be set by the test.
	FlowKeyV flowcontrol.FlowKey

	// AddFunc allows a test to completely override the default Add behavior.
	AddFunc func(item flowcontrol.QueueItemAccessor) error
	// RemoveFunc allows a test to completely override the default Remove behavior.
	RemoveFunc func(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error)
	// CleanupFunc allows a test to completely override the default Cleanup behavior.
	CleanupFunc func(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor
	// DrainFunc allows a test to completely override the default Drain behavior.
	DrainFunc func() []flowcontrol.QueueItemAccessor
	// OrderingPolicyFunc allows a test to override OrderingPolicy.
	OrderingPolicyFunc func() flowcontrol.OrderingPolicy

	// mu protects access to the internal `items` map.
	mu       sync.Mutex
	initOnce sync.Once
	items    map[flowcontrol.QueueItemHandle]flowcontrol.QueueItemAccessor
}

func (m *MockManagedQueue) init() {
	m.initOnce.Do(func() {
		m.items = make(map[flowcontrol.QueueItemHandle]flowcontrol.QueueItemAccessor)
	})
}

// FlowQueueAccessor returns the mock itself, as it fully implements the `flowcontrol.FlowQueueAccessor` interface.
func (m *MockManagedQueue) FlowQueueAccessor() flowcontrol.FlowQueueAccessor {
	return m
}

// Add adds an item to the queue.
// It checks for a test override before locking. If no override is present, it executes the default stateful logic,
// which includes fulfilling the `SafeQueue.Add` contract.
func (m *MockManagedQueue) Add(item flowcontrol.QueueItemAccessor) error {
	// If an override is provided, it is responsible for the full contract, including setting the handle.
	if m.AddFunc != nil {
		return m.AddFunc(item)
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()

	// Fulfill the `SafeQueue.Add` contract: the queue is responsible for setting the handle.
	if item.Handle() == nil {
		item.SetHandle(&mocks.MockQueueItemHandle{})
	}

	m.items[item.Handle()] = item
	return nil
}

// Remove removes an item from the queue. It checks for a test override before locking.
func (m *MockManagedQueue) Remove(handle flowcontrol.QueueItemHandle) (flowcontrol.QueueItemAccessor, error) {
	if m.RemoveFunc != nil {
		return m.RemoveFunc(handle)
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()
	item, ok := m.items[handle]
	if !ok {
		return nil, fmt.Errorf("item with handle %v not found", handle)
	}
	delete(m.items, handle)
	return item, nil
}

// Cleanup removes items matching a predicate. It checks for a test override before locking.
func (m *MockManagedQueue) Cleanup(predicate contracts.PredicateFunc) []flowcontrol.QueueItemAccessor {
	if m.CleanupFunc != nil {
		return m.CleanupFunc(predicate)
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()
	var removed []flowcontrol.QueueItemAccessor
	for handle, item := range m.items {
		if predicate(item) {
			removed = append(removed, item)
			delete(m.items, handle)
		}
	}
	return removed
}

// Drain removes all items from the queue. It checks for a test override before locking.
func (m *MockManagedQueue) Drain() []flowcontrol.QueueItemAccessor {
	if m.DrainFunc != nil {
		return m.DrainFunc()
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()
	drained := make([]flowcontrol.QueueItemAccessor, 0, len(m.items))
	for _, item := range m.items {
		drained = append(drained, item)
	}
	m.items = make(map[flowcontrol.QueueItemHandle]flowcontrol.QueueItemAccessor)
	return drained
}

func (m *MockManagedQueue) FlowKey() flowcontrol.FlowKey { return m.FlowKeyV }
func (m *MockManagedQueue) Name() string                 { return "" }
func (m *MockManagedQueue) OrderingPolicy() flowcontrol.OrderingPolicy {
	if m.OrderingPolicyFunc != nil {
		return m.OrderingPolicyFunc()
	}
	return nil
}

// Len returns the actual number of items currently in the mock queue.
func (m *MockManagedQueue) Len() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()
	return len(m.items)
}

// ByteSize returns the actual total byte size of all items in the mock queue.
func (m *MockManagedQueue) ByteSize() uint64 {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()
	var size uint64
	for _, item := range m.items {
		size += item.OriginalRequest().ByteSize()
	}
	return size
}

// Peek returns the first item found in the mock queue. Note: map iteration order is not guaranteed.
func (m *MockManagedQueue) Peek() flowcontrol.QueueItemAccessor {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.init()
	for _, item := range m.items {
		return item // Return first item found
	}
	return nil // Queue is empty
}