// // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // // 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 tokenizer import ( "strings" "testing" "time" ) type languageDifferentiator struct { input string english string dutch string } // saveEngineType saves the current engineTypeProvider and returns a function // to restore it. Use this when a test modifies the engine type to avoid // leaking global state between tests. func saveEngineType() func() { original := engineTypeProvider return func() { engineTypeProvider = original } } // --------------------------------------------------------------------------- // NumTokensFromString tests // --------------------------------------------------------------------------- func TestNumTokensFromString_Empty(t *testing.T) { if got := NumTokensFromString(""); got != 0 { t.Errorf("expected 0 for empty string, got %d", got) } } func TestNumTokensFromString_Positive(t *testing.T) { for _, s := range []string{"hello world", "你好世界"} { if got := NumTokensFromString(s); got <= 0 { t.Errorf("NumTokensFromString(%q) = %d, want >0", s, got) } } } func TestNumTokensFromString_VariedInputs(t *testing.T) { tests := []struct { name string input string }{ {"ascii letters", "hello world"}, {"chinese characters", "你好世界"}, {"japanese characters", "こんにちは世界"}, {"korean characters", "안녕하세요세계"}, {"emoji", "👋 hello 🌍"}, {"numbers only", "1234567890"}, {"special chars", "a+b=c; d!=e"}, {"newlines and tabs", "line1\nline2\tindented"}, {"mixed content", "RAGFlow 是一款 开源的 RAG (Retrieval-Augmented Generation) 引擎"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := NumTokensFromString(tt.input) if got <= 0 { t.Errorf("NumTokensFromString(%q) = %d, want >0", tt.input, got) } }) } } func TestNumTokensFromString_Consistency(t *testing.T) { inputs := []string{"hello world", "你好世界", "a+b=c; d!=e"} for _, s := range inputs { first := NumTokensFromString(s) second := NumTokensFromString(s) if first != second { t.Errorf("NumTokensFromString(%q) is not consistent: %d vs %d", s, first, second) } } } func TestNumTokensFromString_LongString(t *testing.T) { long := strings.Repeat("the quick brown fox jumps over the lazy dog. ", 200) got := NumTokensFromString(long) if got <= 0 { t.Errorf("NumTokensFromString(long_string) = %d, want >0", got) } } func TestNumTokensFromString_WhitespaceOnly(t *testing.T) { for _, s := range []string{" ", "\t", "\n", " "} { got := NumTokensFromString(s) // Whitespace strings should still produce tokens in BPE encoding if got == 0 { t.Logf("NumTokensFromString(%q) = %d", s, got) } } } // --------------------------------------------------------------------------- // RegisterEngineType tests // --------------------------------------------------------------------------- func TestRegisterEngineType_Basic(t *testing.T) { restore := saveEngineType() defer restore() RegisterEngineType(func() string { return "infinity" }) if got := engineTypeProvider(); got != "infinity" { t.Errorf("expected 'infinity', got %q", got) } } func TestRegisterEngineType_Overwrite(t *testing.T) { restore := saveEngineType() defer restore() RegisterEngineType(func() string { return "first" }) RegisterEngineType(func() string { return "second" }) if got := engineTypeProvider(); got != "second" { t.Errorf("expected 'second', got %q", got) } } // --------------------------------------------------------------------------- // Tokenize tests // --------------------------------------------------------------------------- func TestTokenize_InfinityEngine(t *testing.T) { restore := saveEngineType() defer restore() RegisterEngineType(func() string { return "infinity" }) inputs := []string{"hello world", "你好 世界", "", "a single word"} for _, input := range inputs { got, err := Tokenize(input) if err != nil { t.Errorf("Tokenize(%q) unexpected error: %v", input, err) } if got != input { t.Errorf("Tokenize(%q) = %q, want %q", input, got, input) } } } func TestTokenize_PoolNotInitialized(t *testing.T) { restore := saveEngineType() defer restore() // Ensure engine type is not "infinity" so we hit the pool path RegisterEngineType(func() string { return "" }) _, err := Tokenize("hello world") if err == nil { t.Error("expected error when pool is not initialized, got nil") } } // --------------------------------------------------------------------------- // FineGrainedTokenize tests // --------------------------------------------------------------------------- func TestFineGrainedTokenize_InfinityEngine(t *testing.T) { restore := saveEngineType() defer restore() RegisterEngineType(func() string { return "infinity" }) inputs := []string{"hello world", "测试 分词", ""} for _, input := range inputs { got, err := FineGrainedTokenize(input) if err != nil { t.Errorf("FineGrainedTokenize(%q) unexpected error: %v", input, err) } if got != input { t.Errorf("FineGrainedTokenize(%q) = %q, want %q", input, got, input) } } } func TestFineGrainedTokenize_PoolNotInitialized(t *testing.T) { restore := saveEngineType() defer restore() RegisterEngineType(func() string { return "" }) _, err := FineGrainedTokenize("hello world") if err == nil { t.Error("expected error when pool is not initialized, got nil") } } // --------------------------------------------------------------------------- // Error-path tests for functions that require the pool // --------------------------------------------------------------------------- func TestTokenizeWithPosition_PoolNotInitialized(t *testing.T) { _, err := TokenizeWithPosition("hello world") if err == nil { t.Error("expected error when pool is not initialized, got nil") } } func TestAnalyze_PoolNotInitialized(t *testing.T) { _, err := Analyze("hello world") if err == nil { t.Error("expected error when pool is not initialized, got nil") } } func TestGetTermFreq_PoolNotInitialized(t *testing.T) { got := GetTermFreq("hello") if got != 0 { t.Errorf("expected 0 when pool is not initialized, got %d", got) } } func TestGetTermTag_PoolNotInitialized(t *testing.T) { got := GetTermTag("hello") if got != "" { t.Errorf("expected empty string when pool is not initialized, got %q", got) } } func TestTokenize_DefaultLanguageResetsAnalyzerState(t *testing.T) { restore := saveEngineType() defer restore() RegisterEngineType(func() string { return "" }) if err := Init(&PoolConfig{ DictPath: "", MinSize: 1, MaxSize: 1, IdleTimeout: 5 * time.Second, AcquireTimeout: 5 * time.Second, }); err != nil { t.Fatalf("Failed to initialize pool: %v", err) } defer Close() sample := findEnglishDutchDifferentiator(t) dutchGot, err := New("Dutch").Tokenize(sample.input) if err != nil { t.Fatalf("Tokenize(Dutch, %q) unexpected error: %v", sample.input, err) } if dutchGot != sample.dutch { t.Fatalf("Tokenize(Dutch, %q) = %q, want %q", sample.input, dutchGot, sample.dutch) } defaultGot, err := Tokenize(sample.input) if err != nil { t.Fatalf("Tokenize(default, %q) unexpected error: %v", sample.input, err) } if defaultGot != sample.english { t.Fatalf("Tokenize(default, %q) = %q, want explicit English result %q", sample.input, defaultGot, sample.english) } if defaultGot == dutchGot { t.Fatalf("Tokenize(default, %q) unexpectedly inherited Dutch analyzer state: %q", sample.input, defaultGot) } } func findEnglishDutchDifferentiator(t *testing.T) languageDifferentiator { t.Helper() candidates := []string{ "running", "jumps", "ponies", "studies", "wolves", "relational", "conditionally", } for _, input := range candidates { english, err := New("English").Tokenize(input) if err != nil { t.Fatalf("Tokenize(English, %q) unexpected error: %v", input, err) } dutch, err := New("Dutch").Tokenize(input) if err != nil { t.Fatalf("Tokenize(Dutch, %q) unexpected error: %v", input, err) } if english != dutch { return languageDifferentiator{ input: input, english: english, dutch: dutch, } } } t.Skip("no differentiating tokenizer sample found for English vs Dutch") return languageDifferentiator{} } // --------------------------------------------------------------------------- // Global state tests // --------------------------------------------------------------------------- func TestGetPoolStats_Nil(t *testing.T) { // Note: globalPool is nil by default in unit tests (pool not initialized) stats := GetPoolStats() if stats == nil { t.Fatal("GetPoolStats returned nil") } init, ok := stats["initialized"] if !ok { t.Fatal("missing 'initialized' key") } if init.(bool) { t.Error("expected initialized=false when pool is nil") } } func TestIsInitialized_Default(t *testing.T) { if IsInitialized() { t.Error("expected IsInitialized() = false when pool is not initialized") } } func TestClose_Nil(t *testing.T) { // Close should be safe to call with nil globalPool Close() // no panic = pass } func TestClose_NilGlobalPool(t *testing.T) { // Call Close directly after ensuring globalPool is nil // (concurrent test may have initialized it, so handle gracefully) defer func() { if r := recover(); r != nil { t.Errorf("Close() panicked: %v", r) } }() Close() }