From a28f2f9c86828e2533fcf516fb0b8b3577ec3f6f Mon Sep 17 00:00:00 2001
From: yuluo-yx <yuluo08290126@gmail.com>
Date: Sat, 6 Sep 2025 16:54:34 +0800
Subject: [PATCH] feat: optimize classifier and add ttft unit test

Signed-off-by: yuluo-yx <yuluo08290126@gmail.com>
---
 .../pkg/utils/ttft/calculator_test.go         | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 src/semantic-router/pkg/utils/ttft/calculator_test.go

diff --git a/src/semantic-router/pkg/utils/ttft/calculator_test.go b/src/semantic-router/pkg/utils/ttft/calculator_test.go
new file mode 100644
index 00000000..bf4b3fa1
--- /dev/null
+++ b/src/semantic-router/pkg/utils/ttft/calculator_test.go
@@ -0,0 +1,56 @@
+package ttft
+
+import (
+	"testing"
+
+	"github.com/vllm-project/semantic-router/semantic-router/pkg/config"
+)
+
+func TestComputeBaseTTFT(t *testing.T) {
+
+	gpuConfig := config.GPUConfig{
+		FLOPS: 1e12, // 1 TFLOP
+		HBM:   1e11, // 100 GB/s
+	}
+	calculator := NewCalculator(gpuConfig)
+
+	routerCfg := &config.RouterConfig{}
+	// Mock config methods if needed, or set up fields so that
+	// GetModelParamCount, GetModelBatchSize, GetModelContextSize return defaults
+
+	ttft := calculator.ComputeBaseTTFT("test-model", routerCfg)
+	if ttft <= 0 {
+		t.Errorf("Expected TTFT > 0, got %f", ttft)
+	}
+}
+
+func TestInitializeModelTTFT(t *testing.T) {
+	gpuConfig := config.GPUConfig{
+		FLOPS: 1e12,
+		HBM:   1e11,
+	}
+	calculator := NewCalculator(gpuConfig)
+
+	// Minimal mock config with two categories and models
+	routerCfg := &config.RouterConfig{
+		Categories: []config.Category{
+			{
+				ModelScores: []config.ModelScore{
+					{Model: "model-a", Score: 0.9},
+					{Model: "model-b", Score: 0.8},
+				},
+			},
+		},
+		DefaultModel: "model-default",
+	}
+
+	modelTTFT := calculator.InitializeModelTTFT(routerCfg)
+	if len(modelTTFT) != 3 {
+		t.Errorf("Expected 3 models in TTFT map, got %d", len(modelTTFT))
+	}
+	for model, ttft := range modelTTFT {
+		if ttft <= 0 {
+			t.Errorf("Model %s has non-positive TTFT: %f", model, ttft)
+		}
+	}
+}