update pool tests

Signed-off-by: Sage Ahrac <[email protected]>

Author: sagiahrac

pkg/tokenization/pool_test.go

@@ -100,12 +100,12 @@ func TestPool_ProcessTask(t *testing.T) {
 	expectedOffsets := []tokenizers.Offset{{0, 5}, {6, 11}}
 
 	// Mock FindLongestContainedTokens to return low overlap ratio
-	mockIndexer.On("FindLongestContainedTokens", task.Prompt, task.ModelName).Return([]uint32{}, 0.0)
+	mockIndexer.On("FindLongestContainedTokens", task.Prompt).Return([]uint32{}, 0.0)
 
 	mockTokenizer.On("Encode", task.Prompt, task.ModelName).Return(expectedTokens, expectedOffsets, nil)
 
 	// Verify that indexer receives exactly the same tokens and offsets that tokenizer returned
-	mockIndexer.On("AddTokenization", task.ModelName, task.Prompt, expectedTokens, expectedOffsets).Return(nil)
+	mockIndexer.On("AddTokenization", task.Prompt, expectedTokens, expectedOffsets).Return(nil)
 
 	// Execute
 	err := pool.processTask(task)
@@ -127,8 +127,8 @@ func TestPool_RunIntegration(t *testing.T) {
 
 	// Setup mock expectations for each prompt
 	for _, prompt := range prompts {
-		mockIndexer.On("FindLongestContainedTokens", prompt, testModelName).Return([]uint32{}, 0.0)
-		mockIndexer.On("AddTokenization", testModelName, prompt,
+		mockIndexer.On("FindLongestContainedTokens", prompt).Return([]uint32{}, 0.0)
+		mockIndexer.On("AddTokenization", prompt,
 			mock.Anything, mock.Anything).Return(nil).Once()
 	}