From 0f259c18dd75e59fa0fb243498eddd81532a9bc0 Mon Sep 17 00:00:00 2001
From: xgopilot <noreply@goplus.org>
Date: Thu, 23 Apr 2026 05:00:14 +0000
Subject: [PATCH] test(provider,scripts): improve coverage for estimator and
 migration output

Generated with [codeagent](https://github.com/qbox/codeagent)
Co-authored-by: phantom5099 <245659304+phantom5099@users.noreply.github.com>
---
 internal/provider/estimate_test.go          | 71 ++++++++++++++++
 scripts/migrate_context_budget/main_test.go | 90 +++++++++++++++++++++
 2 files changed, 161 insertions(+)
 create mode 100644 internal/provider/estimate_test.go

diff --git a/internal/provider/estimate_test.go b/internal/provider/estimate_test.go
new file mode 100644
index 00000000..5fecb135
--- /dev/null
+++ b/internal/provider/estimate_test.go
@@ -0,0 +1,71 @@
+package provider
+
+import (
+	"testing"
+
+	providertypes "neo-code/internal/provider/types"
+)
+
+func TestEstimateSerializedPayloadTokens(t *testing.T) {
+	t.Parallel()
+
+	tokens, err := EstimateSerializedPayloadTokens(map[string]any{
+		"model": "x",
+		"input": "hello",
+	})
+	if err != nil {
+		t.Fatalf("EstimateSerializedPayloadTokens() error = %v", err)
+	}
+	if tokens <= 0 {
+		t.Fatalf("EstimateSerializedPayloadTokens() = %d, want > 0", tokens)
+	}
+}
+
+func TestEstimateSerializedPayloadTokensMarshalError(t *testing.T) {
+	t.Parallel()
+
+	if _, err := EstimateSerializedPayloadTokens(make(chan int)); err == nil {
+		t.Fatal("EstimateSerializedPayloadTokens() expected marshal error, got nil")
+	}
+}
+
+func TestEstimateTextTokens(t *testing.T) {
+	t.Parallel()
+
+	if got := EstimateTextTokens(""); got != 0 {
+		t.Fatalf("EstimateTextTokens(\"\") = %d, want 0", got)
+	}
+	if got := EstimateTextTokens("1234"); got != 2 {
+		t.Fatalf("EstimateTextTokens(\"1234\") = %d, want 2", got)
+	}
+}
+
+func TestBuildGenerateRequestSignature(t *testing.T) {
+	t.Parallel()
+
+	reqA := providertypes.GenerateRequest{
+		Model: "gpt",
+		Messages: []providertypes.Message{
+			{
+				Role:  providertypes.RoleUser,
+				Parts: []providertypes.ContentPart{providertypes.NewTextPart("hello")},
+			},
+		},
+	}
+	reqB := reqA
+	reqC := reqA
+	reqC.Model = "gpt-2"
+
+	sigA := BuildGenerateRequestSignature(reqA)
+	sigB := BuildGenerateRequestSignature(reqB)
+	sigC := BuildGenerateRequestSignature(reqC)
+	if sigA == "" {
+		t.Fatal("BuildGenerateRequestSignature(reqA) returned empty signature")
+	}
+	if sigA != sigB {
+		t.Fatalf("same request should have same signature: %q != %q", sigA, sigB)
+	}
+	if sigA == sigC {
+		t.Fatalf("different requests should have different signatures: %q == %q", sigA, sigC)
+	}
+}
diff --git a/scripts/migrate_context_budget/main_test.go b/scripts/migrate_context_budget/main_test.go
index dff132e4..7c2a37ef 100644
--- a/scripts/migrate_context_budget/main_test.go
+++ b/scripts/migrate_context_budget/main_test.go
@@ -1,7 +1,14 @@
 package main
 
 import (
+	"bytes"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
 	"testing"
+
+	"neo-code/internal/config"
 )
 
 func TestDefaultBaseDirReturnsPath(t *testing.T) {
@@ -11,3 +18,86 @@ func TestDefaultBaseDirReturnsPath(t *testing.T) {
 		t.Fatal("expected non-empty default base dir")
 	}
 }
+
+func TestPrintMigrationResultChangedDryRun(t *testing.T) {
+	output := captureStdout(t, func() {
+		printMigrationResult(config.ContextBudgetMigrationResult{
+			Path:    "/tmp/config.yaml",
+			Changed: true,
+		}, true)
+	})
+
+	if !strings.Contains(output, "[DRY-RUN] 将迁移 /tmp/config.yaml") {
+		t.Fatalf("unexpected output: %q", output)
+	}
+}
+
+func TestPrintMigrationResultChangedWithBackup(t *testing.T) {
+	output := captureStdout(t, func() {
+		printMigrationResult(config.ContextBudgetMigrationResult{
+			Path:    "/tmp/config.yaml",
+			Changed: true,
+			Backup:  "/tmp/config.yaml.bak",
+		}, false)
+	})
+
+	if !strings.Contains(output, "已迁移 /tmp/config.yaml (备份: /tmp/config.yaml.bak)") {
+		t.Fatalf("unexpected output: %q", output)
+	}
+}
+
+func TestPrintMigrationResultNotChangedWithNotes(t *testing.T) {
+	output := captureStdout(t, func() {
+		printMigrationResult(config.ContextBudgetMigrationResult{
+			Path:   "/tmp/config.yaml",
+			Reason: "未检测到 context.auto_compact",
+			Notes:  []string{" note-a ", "note-b"},
+		}, false)
+	})
+
+	if !strings.Contains(output, "说明: note-a") {
+		t.Fatalf("missing note-a in output: %q", output)
+	}
+	if !strings.Contains(output, "说明: note-b") {
+		t.Fatalf("missing note-b in output: %q", output)
+	}
+	if !strings.Contains(output, "跳过: /tmp/config.yaml (未检测到 context.auto_compact)") {
+		t.Fatalf("missing skip line in output: %q", output)
+	}
+}
+
+func TestDefaultBaseDirUsesHome(t *testing.T) {
+	tempHome := t.TempDir()
+	t.Setenv("HOME", tempHome)
+	want := filepath.Join(tempHome, ".neocode")
+	if got := defaultBaseDir(); got != want {
+		t.Fatalf("defaultBaseDir() = %q, want %q", got, want)
+	}
+}
+
+func captureStdout(t *testing.T, fn func()) string {
+	t.Helper()
+
+	originalStdout := os.Stdout
+	reader, writer, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("os.Pipe() error = %v", err)
+	}
+	os.Stdout = writer
+	defer func() {
+		os.Stdout = originalStdout
+	}()
+
+	done := make(chan string, 1)
+	go func() {
+		var buf bytes.Buffer
+		_, _ = io.Copy(&buf, reader)
+		done <- buf.String()
+	}()
+
+	fn()
+	_ = writer.Close()
+	output := <-done
+	_ = reader.Close()
+	return output
+}