pool_test.go

/*
Copyright 2025 The llm-d 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.
*/

//nolint:testpackage // need to test internal types
package tokenization

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"k8s.io/client-go/util/workqueue"

	types "github.com/llm-d/llm-d-kv-cache/pkg/tokenization/types"
)

// MockTokenizer implements the Tokenizer interface for testing.
type MockTokenizer struct {
	mock.Mock
}

func (m *MockTokenizer) RenderChat(renderReq *types.RenderChatRequest) ([]uint32, []types.Offset, error) {
	args := m.Called(renderReq)
	tokenIface := args.Get(0)
	if tokenIface == nil {
		return nil, nil, args.Error(2)
	}
	tokens, ok := tokenIface.([]uint32)
	if !ok {
		panic("MockTokenizer.RenderChat: expected []uint32 from mock, got unexpected type")
	}
	offsetIface := args.Get(1)
	if offsetIface == nil {
		return nil, nil, args.Error(2)
	}
	offsets, ok := offsetIface.([]types.Offset)
	if !ok {
		panic("MockTokenizer.RenderChat: expected []types.Offset from mock, got unexpected type")
	}
	return tokens, offsets, args.Error(2)
}

func (m *MockTokenizer) Render(prompt string) ([]uint32, []types.Offset, error) {
	args := m.Called(prompt)
	tokenIface := args.Get(0)
	if tokenIface == nil {
		return nil, nil, args.Error(2)
	}
	tokens, ok := tokenIface.([]uint32)
	if !ok {
		panic("MockTokenizer.Render: expected []uint32 from mock, got unexpected type")
	}
	offsetIface := args.Get(1)
	if offsetIface == nil {
		return nil, nil, args.Error(2)
	}
	offsets, ok := offsetIface.([]types.Offset)
	if !ok {
		panic("MockTokenizer.Render: expected []types.Offset from mock, got unexpected type")
	}
	return tokens, offsets, args.Error(2)
}

func (m *MockTokenizer) Close() error {
	return nil
}

func (m *MockTokenizer) Type() string {
	return "mock"
}

func TestPool_ProcessTask(t *testing.T) {
	mockTokenizer := &MockTokenizer{}

	pool := &Pool{
		modelName: "test-model",
		workers:   1,
		tokenizer: mockTokenizer,
	}

	task := Task{
		Prompt: "hello world",
	}

	// Setup specific mock return values
	expectedTokens := []uint32{12345, 67890, 11111}
	expectedOffsets := []types.Offset{{0, 5}, {6, 11}}

	mockTokenizer.On("Render", task.Prompt).
		Return(expectedTokens, expectedOffsets, nil)

	// Execute
	err := pool.processTask(task)

	// Assert
	assert.NoError(t, err)
	mockTokenizer.AssertExpectations(t)
}

func TestPool_WorkerLoop(t *testing.T) {
	specs := map[string]struct {
		setupMocks func(*MockTokenizer)
		genTasks   func() ([]Task, chan tokenizationResponse)
		verify     func(t *testing.T, pool *Pool, tasks []Task, resultChan chan tokenizationResponse)
	}{
		"successful task processing": {
			setupMocks: func(mt *MockTokenizer) {
				mt.On("Render", "test prompt").
					Return([]uint32{1, 2, 3}, []types.Offset{{0, 4}}, nil)
			},
			genTasks: func() ([]Task, chan tokenizationResponse) {
				return []Task{{Prompt: "test prompt"}}, nil
			},
			verify: func(t *testing.T, pool *Pool, tasks []Task, resultChan chan tokenizationResponse) {}, //nolint:thelper // noop
		},
		"task with result channel": {
			setupMocks: func(mt *MockTokenizer) {
				mt.On("Render", "test with channel").
					Return([]uint32{10, 20, 30}, []types.Offset{{0, 4}}, nil)
			},
			genTasks: func() ([]Task, chan tokenizationResponse) {
				ch := make(chan tokenizationResponse, 1)
				return []Task{{
					Prompt:   "test with channel",
					ResultCh: ch,
				}}, ch
			},
			verify: func(t *testing.T, pool *Pool, tasks []Task, resultCh chan tokenizationResponse) {
				t.Helper()
				require.Eventually(t, func() bool {
					if result, ok := <-resultCh; ok {
						assert.Equal(t, []uint32{10, 20, 30}, result.Tokens)
						return true
					}
					return false
				}, time.Second, 10*time.Millisecond)

				// Verify channel is closed
				require.Eventually(t, func() bool {
					_, ok := <-resultCh
					return !ok
				}, time.Second, 10*time.Millisecond)
			},
		},
		"multiple tasks processing": {
			setupMocks: func(mt *MockTokenizer) {
				for i := range 5 {
					prompt := "prompt " + string(rune('a'+i))
					tokens := []uint32{uint32(i), uint32(i + 1)} // i is bounded by range 5, no overflow
					offsets := []types.Offset{{0, 6}}

					mt.On("Render", prompt).
						Return(tokens, offsets, nil).Once()
				}
			},
			genTasks: func() ([]Task, chan tokenizationResponse) {
				tasks := make([]Task, 5)
				for i := range 5 {
					tasks[i] = Task{Prompt: "prompt " + string(rune('a'+i))}
				}
				return tasks, nil
			},
			verify: func(t *testing.T, pool *Pool, tasks []Task, resultChan chan tokenizationResponse) {
				t.Helper()
				require.Eventually(t, func() bool {
					return pool.queue.Len() == 0
				}, time.Second, 10*time.Millisecond, "queue should be drained")
			},
		},
		"max retries exceeded": {
			setupMocks: func(mt *MockTokenizer) {
				// Mock will fail every time, causing retries
				mt.On("Render", "failing prompt").Return(
					[]uint32{}, []types.Offset{}, assert.AnError)
			},
			genTasks: func() ([]Task, chan tokenizationResponse) {
				ch := make(chan tokenizationResponse, 1)
				return []Task{{
					Prompt:   "failing prompt",
					ResultCh: ch,
				}}, ch
			},
			verify: func(t *testing.T, pool *Pool, tasks []Task, resultCh chan tokenizationResponse) {
				t.Helper()
				require.Eventually(t, func() bool { // channel is closed, when max retries exceeded
					if result, ok := <-resultCh; !ok {
						assert.Equal(t, tokenizationResponse{}, result)
						return true
					}
					return false
				}, time.Second, 10*time.Millisecond)

				require.Eventually(t, func() bool {
					return pool.queue.Len() == 0
				}, time.Second, 10*time.Millisecond)
			},
		},
	}
	for name, tt := range specs {
		t.Run(name, func(t *testing.T) {
			mockTokenizer := &MockTokenizer{}

			tt.setupMocks(mockTokenizer)
			pool := &Pool{
				modelName: "test-model",
				workers:   1,
				queue:     workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[Task]()),
				tokenizer: mockTokenizer,
			}

			tasks, resultChan := tt.genTasks()
			for _, task := range tasks {
				pool.queue.Add(task)
			}

			pool.wg.Add(1)
			go pool.workerLoop(0)

			tt.verify(t, pool, tasks, resultChan)

			// Shutdown
			pool.queue.ShutDown()
			pool.wg.Wait()

			// Assert expectations
			mockTokenizer.AssertExpectations(t)
		})
	}
}