xinmi/scripts/graph_extraction_probe/main_test.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
)

type fakeProbeChat struct {
	chatDelegate
	response *types.ChatResponse
	err      error
}

func (f *fakeProbeChat) Chat(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error) {
	return f.response, f.err
}

func TestRecordingChatPreservesPerCallDiagnosticsWithoutContent(t *testing.T) {
	fake := &fakeProbeChat{response: &types.ChatResponse{
		Content: "answer", ReasoningContent: "private reasoning", FinishReason: "length",
		Usage: types.TokenUsage{PromptTokens: 10, CompletionTokens: 4, TotalTokens: 14},
	}}
	recorder := &recordingChat{chatDelegate: fake}
	thinking := false
	options := &chat.ChatOptions{MaxTokens: 4096, Thinking: &thinking}
	first, err := recorder.Chat(context.Background(), nil, options)
	if err != nil || first != fake.response {
		t.Fatal("recorder changed the response")
	}
	thinking = true
	fake.response = &types.ChatResponse{Content: "\u4e2d\u6587", FinishReason: "stop"}
	if _, err := recorder.Chat(context.Background(), nil, options); err != nil {
		t.Fatal(err)
	}
	if recorder.calls != 2 || len(recorder.records) != 2 {
		t.Fatal("missing call diagnostics")
	}
	a, b := recorder.records[0], recorder.records[1]
	if a.Number != 1 || a.MaxTokens != 4096 || a.Thinking == nil || *a.Thinking ||
		!a.ResponseReceived || a.Failed || a.FinishReason != "length" ||
		a.ContentRunes != 6 || a.ReasoningRunes != 17 || a.Usage.TotalTokens != 14 {
		t.Fatalf("unexpected first call metadata: %+v", a)
	}
	if b.Number != 2 || b.Thinking == nil || !*b.Thinking || b.ContentRunes != 2 || b.ReasoningRunes != 0 {
		t.Fatalf("unexpected second call metadata: %+v", b)
	}
	encoded, err := json.Marshal(recorder.records)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(encoded), "private reasoning") || strings.Contains(string(encoded), "answer") {
		t.Fatal("diagnostics exposed response content")
	}
}

func TestRecordingChatRecordsNoResponseWithoutLeakingError(t *testing.T) {
	want := errors.New("private upstream error")
	recorder := &recordingChat{chatDelegate: &fakeProbeChat{err: want}}
	response, err := recorder.Chat(context.Background(), nil, nil)
	if response != nil || err != want || len(recorder.records) != 1 {
		t.Fatal("recorder changed the failure")
	}
	record := recorder.records[0]
	if record.ResponseReceived || !record.Failed || record.Thinking != nil || record.ContentRunes != 0 || record.ReasoningRunes != 0 {
		t.Fatalf("unexpected failure metadata: %+v", record)
	}
	encoded, _ := json.Marshal(record)
	if strings.Contains(string(encoded), want.Error()) {
		t.Fatal("diagnostics exposed upstream error")
	}
}