Enterprise-Knowledge-Base/cli/cmd/chat/chat_test.go

package chat

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

	"github.com/Tencent/WeKnora/cli/internal/cmdutil"
	"github.com/Tencent/WeKnora/cli/internal/iostreams"
	sdk "github.com/Tencent/WeKnora/client"
)

// fakeChatService implements ChatService for unit tests. Tests configure the
// callback driver via streamEvents (delivered in order) and observe captured
// inputs through the exported fields.
type fakeChatService struct {
	createSessionResp *sdk.Session
	createSessionErr  error
	createCalled      bool

	streamErr      error
	streamEvents   []*sdk.StreamResponse
	gotSessionID   string
	gotRequest     *sdk.KnowledgeQARequest
	streamCalled   bool
	cbReturnsError error // if set, callback aborts after first event with this error
}

func (f *fakeChatService) CreateSession(_ context.Context, req *sdk.CreateSessionRequest) (*sdk.Session, error) {
	f.createCalled = true
	if f.createSessionErr != nil {
		return nil, f.createSessionErr
	}
	if f.createSessionResp != nil {
		return f.createSessionResp, nil
	}
	// Default: return a deterministic session id derived from the title so
	// JSON assertions don't depend on uuid generation.
	return &sdk.Session{ID: "sess_auto", Title: req.Title}, nil
}

func (f *fakeChatService) KnowledgeQAStream(ctx context.Context, sessionID string, req *sdk.KnowledgeQARequest, cb func(*sdk.StreamResponse) error) error {
	f.streamCalled = true
	f.gotSessionID = sessionID
	f.gotRequest = req
	for _, ev := range f.streamEvents {
		if err := cb(ev); err != nil {
			return err
		}
		if f.cbReturnsError != nil {
			return f.cbReturnsError
		}
	}
	return f.streamErr
}

// Sanity: fakeChatService must satisfy ChatService. Mirrors the production
// var _ ChatService = (*sdk.Client)(nil) check at the bottom of chat.go.
var _ ChatService = (*fakeChatService)(nil)

// textOpts returns a FormatOptions configured for the text render path —
// the most common shape under test.
func textOpts() *cmdutil.FormatOptions {
	return &cmdutil.FormatOptions{Mode: cmdutil.FormatText}
}

func TestChat_StreamMode(t *testing.T) {
	out, errBuf := iostreams.SetForTestWithTTY(t)
	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{
			{ResponseType: sdk.ResponseTypeAnswer, Content: "Hello "},
			{ResponseType: sdk.ResponseTypeAnswer, Content: "world"},
			{ResponseType: sdk.ResponseTypeReferences, KnowledgeReferences: []*sdk.SearchResult{
				{KnowledgeID: "k1", KnowledgeTitle: "Doc One", Score: 0.42},
			}},
			{ResponseType: sdk.ResponseTypeComplete, Done: true},
		},
	}
	opts := &Options{Query: "hi", KBID: "kb_1"}
	if err := runChat(context.Background(), opts, textOpts(), svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}
	got := out.String()
	if !strings.Contains(got, "Hello world") {
		t.Errorf("stdout missing streamed content: %q", got)
	}
	if !strings.Contains(got, "References") {
		t.Errorf("stdout missing references footer: %q", got)
	}
	if !strings.Contains(got, "Doc One") {
		t.Errorf("references should render KnowledgeTitle, got %q", got)
	}
	// auto-created session must announce itself on stderr
	if !strings.Contains(errBuf.String(), "session: sess_auto") {
		t.Errorf("expected stderr session hint, got %q", errBuf.String())
	}
	if !svc.createCalled {
		t.Error("expected CreateSession invocation when SessionID empty")
	}
	if svc.gotSessionID != "sess_auto" {
		t.Errorf("stream sessionID: got %q want sess_auto", svc.gotSessionID)
	}
	if svc.gotRequest == nil || svc.gotRequest.Channel != "api" {
		t.Errorf("expected Channel=api, got %+v", svc.gotRequest)
	}
}

// TestChat_NDJSON_FirstLineIsInit verifies that the NDJSON path (--format json)
// always injects an "init" line first carrying session_id and kb_id.
func TestChat_NDJSON_FirstLineIsInit(t *testing.T) {
	out, errBuf := iostreams.SetForTest(t)

	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{
			{ResponseType: sdk.ResponseTypeAnswer, Content: "answer"},
			{ResponseType: sdk.ResponseTypeComplete, Done: true},
		},
	}
	opts := &Options{Query: "q", KBID: "kb_42"}
	fopts := &cmdutil.FormatOptions{Mode: cmdutil.FormatJSON}
	if err := runChat(context.Background(), opts, fopts, svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}

	// NDJSON mode must NOT print the session hint to stderr.
	if errBuf.Len() != 0 {
		t.Errorf("expected empty stderr in NDJSON mode, got %q", errBuf.String())
	}

	lines := strings.Split(strings.TrimRight(out.String(), "\n"), "\n")
	if len(lines) == 0 {
		t.Fatal("no output")
	}
	var first struct {
		Type      string `json:"type"`
		SessionID string `json:"session_id"`
		KBID      string `json:"kb_id"`
	}
	if err := json.Unmarshal([]byte(lines[0]), &first); err != nil {
		t.Fatalf("first line not JSON: %v\n  %s", err, lines[0])
	}
	if first.Type != "init" {
		t.Errorf("first line type: got %q, want init", first.Type)
	}
	if first.SessionID != "sess_auto" {
		t.Errorf("init.session_id: got %q, want sess_auto", first.SessionID)
	}
	if first.KBID != "kb_42" {
		t.Errorf("init.kb_id: got %q, want kb_42", first.KBID)
	}
}

// TestChat_NDJSON_PassthroughEvents verifies that the NDJSON path emits
// init + N SDK events = N+1 total lines (no buffering, no extra wrapping).
func TestChat_NDJSON_PassthroughEvents(t *testing.T) {
	out, _ := iostreams.SetForTest(t)

	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{
			{ResponseType: sdk.ResponseTypeAnswer, Content: "hello"},
			{ResponseType: sdk.ResponseTypeComplete, Done: true},
		},
	}
	opts := &Options{Query: "q", KBID: "kb_x"}
	fopts := &cmdutil.FormatOptions{Mode: cmdutil.FormatJSON}
	if err := runChat(context.Background(), opts, fopts, svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}

	lines := strings.Split(strings.TrimRight(out.String(), "\n"), "\n")
	// 1 init + 2 SDK events = 3 lines.
	if len(lines) != 3 {
		t.Fatalf("got %d lines, want 3:\n%s", len(lines), out.String())
	}
	// Each must be valid JSON.
	for i, line := range lines {
		var obj map[string]any
		if err := json.Unmarshal([]byte(line), &obj); err != nil {
			t.Errorf("line %d not valid JSON: %v\n  %s", i+1, err, line)
		}
	}
}

// TestChat_NDJSON_JSONEqualsNDJSON verifies that --format json and --format ndjson
// produce identical byte output on a streaming command.
func TestChat_NDJSON_JSONEqualsNDJSON(t *testing.T) {
	events := []*sdk.StreamResponse{
		{ResponseType: sdk.ResponseTypeAnswer, Content: "hello"},
		{ResponseType: sdk.ResponseTypeComplete, Done: true},
	}

	runWith := func(mode cmdutil.FormatMode) string {
		out, _ := iostreams.SetForTest(t)
		svc := &fakeChatService{streamEvents: events}
		opts := &Options{Query: "q", KBID: "kb_x"}
		fopts := &cmdutil.FormatOptions{Mode: mode}
		if err := runChat(context.Background(), opts, fopts, svc); err != nil {
			t.Fatalf("runChat(%s): %v", mode, err)
		}
		return out.String()
	}

	jsonOut := runWith(cmdutil.FormatJSON)
	ndjsonOut := runWith(cmdutil.FormatNDJSON)
	if jsonOut != ndjsonOut {
		t.Errorf("--format json and --format ndjson differ:\n  json:   %q\n  ndjson: %q", jsonOut, ndjsonOut)
	}
}

func TestChat_NonTTY_AccumulateMode(t *testing.T) {
	// Non-TTY iostreams forces accumulate mode.
	out, _ := iostreams.SetForTest(t)
	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{
			{ResponseType: sdk.ResponseTypeAnswer, Content: "piped"},
			{ResponseType: sdk.ResponseTypeComplete, Done: true},
		},
	}
	opts := &Options{Query: "q", KBID: "kb"}
	if err := runChat(context.Background(), opts, textOpts(), svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}
	if !strings.Contains(out.String(), "piped") {
		t.Errorf("expected accumulated answer, got %q", out.String())
	}
}

func TestChat_SessionIDProvided(t *testing.T) {
	_, errBuf := iostreams.SetForTestWithTTY(t)
	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{{ResponseType: sdk.ResponseTypeComplete, Done: true}},
	}
	opts := &Options{Query: "q", KBID: "kb", SessionID: "sess_existing"}
	if err := runChat(context.Background(), opts, textOpts(), svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}
	if svc.createCalled {
		t.Error("CreateSession must NOT be invoked when --session is provided")
	}
	if svc.gotSessionID != "sess_existing" {
		t.Errorf("stream sessionID: got %q want sess_existing", svc.gotSessionID)
	}
	// No auto-create message because the user supplied the id.
	if strings.Contains(errBuf.String(), "session:") {
		t.Errorf("unexpected session hint emitted with explicit --session: %q", errBuf.String())
	}
}

func TestChat_KBIDRequired(t *testing.T) {
	_, _ = iostreams.SetForTest(t)
	svc := &fakeChatService{}
	// Run with KBID empty (bypassing the cobra resolver).
	opts := &Options{Query: "q"}
	err := runChat(context.Background(), opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeKBIDRequired {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeKBIDRequired)
	}
	if svc.createCalled || svc.streamCalled {
		t.Error("KB validation must short-circuit before any SDK call")
	}
}

func TestChat_EmptyQuery(t *testing.T) {
	_, _ = iostreams.SetForTest(t)
	svc := &fakeChatService{}
	opts := &Options{Query: "", KBID: "kb"}
	err := runChat(context.Background(), opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeInputInvalidArgument {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeInputInvalidArgument)
	}
}

func TestChat_SDKError_PreStream(t *testing.T) {
	// SDK fails before any event arrives → ClassifyHTTPError mapping.
	// "HTTP error 401: ..." → auth.unauthenticated.
	_, _ = iostreams.SetForTest(t)
	svc := &fakeChatService{
		streamErr: errors.New("HTTP error 401: token rejected"),
	}
	opts := &Options{Query: "q", KBID: "kb"}
	err := runChat(context.Background(), opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeAuthUnauthenticated {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeAuthUnauthenticated)
	}
}

func TestChat_SDKError_MidStream_AbortsAsSSE(t *testing.T) {
	// Some content arrived, then the stream errored without a Done event →
	// CodeSSEStreamAborted (separate from generic transport failure).
	_, _ = iostreams.SetForTest(t)
	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{{Content: "partial"}},
		streamErr:    errors.New("connection reset"),
	}
	opts := &Options{Query: "q", KBID: "kb"}
	err := runChat(context.Background(), opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeSSEStreamAborted {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeSSEStreamAborted)
	}
}

func TestChat_ContextCancelled(t *testing.T) {
	_, _ = iostreams.SetForTest(t)
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // simulate Ctrl-C delivered before the SDK returns.
	svc := &fakeChatService{streamErr: context.Canceled}
	opts := &Options{Query: "q", KBID: "kb"}
	err := runChat(ctx, opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeOperationCancelled {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeOperationCancelled)
	}
}

func TestChat_SessionCreateFails(t *testing.T) {
	_, _ = iostreams.SetForTest(t)
	svc := &fakeChatService{
		createSessionErr: errors.New("dial tcp: connection refused"),
	}
	opts := &Options{Query: "q", KBID: "kb"}
	err := runChat(context.Background(), opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeSessionCreateFailed {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeSessionCreateFailed)
	}
	if svc.streamCalled {
		t.Error("stream must not be invoked after session creation failed")
	}
}

func TestChat_SessionCreate404SurfacesNotFound(t *testing.T) {
	// HTTP-shaped session-create failures should NOT collapse into the
	// session_create_failed bucket; they keep their canonical mapping so
	// agents can react to e.g. resource.not_found.
	_, _ = iostreams.SetForTest(t)
	svc := &fakeChatService{
		createSessionErr: errors.New("HTTP error 404: tenant not found"),
	}
	opts := &Options{Query: "q", KBID: "kb"}
	err := runChat(context.Background(), opts, textOpts(), svc)
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	var typed *cmdutil.Error
	if !errors.As(err, &typed) {
		t.Fatalf("expected *cmdutil.Error, got %T", err)
	}
	if typed.Code != cmdutil.CodeResourceNotFound {
		t.Errorf("code: got %q want %q", typed.Code, cmdutil.CodeResourceNotFound)
	}
}

// TestChat_NDJSON_InitIncludesProfile verifies that when a profile is set,
// the NDJSON init event carries the profile field.
func TestChat_NDJSON_InitIncludesProfile(t *testing.T) {
	cmdutil.SetProfile("prod")
	t.Cleanup(func() { cmdutil.SetProfile("") })

	out, _ := iostreams.SetForTest(t)
	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{
			{ResponseType: sdk.ResponseTypeComplete, Done: true},
		},
	}
	opts := &Options{Query: "q", KBID: "kb_x"}
	fopts := &cmdutil.FormatOptions{Mode: cmdutil.FormatNDJSON}
	if err := runChat(context.Background(), opts, fopts, svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}

	lines := strings.Split(strings.TrimRight(out.String(), "\n"), "\n")
	if len(lines) == 0 {
		t.Fatal("no output")
	}
	var initLine struct {
		Type    string `json:"type"`
		Profile string `json:"profile"`
	}
	if err := json.Unmarshal([]byte(lines[0]), &initLine); err != nil {
		t.Fatalf("first line not JSON: %v\n  %s", err, lines[0])
	}
	if initLine.Type != "init" {
		t.Errorf("type: got %q, want init", initLine.Type)
	}
	if initLine.Profile != "prod" {
		t.Errorf("profile: got %q, want prod", initLine.Profile)
	}
}

func TestChat_FormatNDJSON_PassthroughsSDKEvents(t *testing.T) {
	// Fake stream emits 3 events: thinking, answer, complete.
	// With the init injection, total output is 4 lines (1 init + 3 SDK events).
	svc := &fakeChatService{
		streamEvents: []*sdk.StreamResponse{
			{ResponseType: sdk.ResponseTypeThinking, Content: "search KB"},
			{ResponseType: sdk.ResponseTypeAnswer, Content: "hello"},
			{ResponseType: sdk.ResponseTypeComplete, Done: true, SessionID: "sess_x"},
		},
	}
	out, _ := iostreams.SetForTest(t)

	opts := &Options{Query: "hi", KBID: "kb_x"}
	fopts := &cmdutil.FormatOptions{Mode: cmdutil.FormatNDJSON}
	if err := runChat(context.Background(), opts, fopts, svc); err != nil {
		t.Fatalf("runChat: %v", err)
	}
	lines := strings.Split(strings.TrimRight(out.String(), "\n"), "\n")
	// 1 init + 3 SDK events = 4 lines.
	if len(lines) != 4 {
		t.Fatalf("got %d lines, want 4:\n%s", len(lines), out.String())
	}

	// First line: CLI-injected init event.
	var initLine map[string]any
	if err := json.Unmarshal([]byte(lines[0]), &initLine); err != nil {
		t.Fatalf("line 1 (init) not JSON: %v", err)
	}
	if initLine["type"] != "init" {
		t.Errorf("first line type=%v, want init", initLine["type"])
	}

	// Second line: thinking event (SDK passthrough).
	var second map[string]any
	if err := json.Unmarshal([]byte(lines[1]), &second); err != nil {
		t.Fatalf("line 2 not JSON: %v", err)
	}
	if second["response_type"] != "thinking" {
		t.Errorf("second event response_type=%v, want thinking", second["response_type"])
	}
	// Third line: answer.
	var third map[string]any
	if err := json.Unmarshal([]byte(lines[2]), &third); err != nil {
		t.Fatalf("line 3 not JSON: %v", err)
	}
	if third["response_type"] != "answer" {
		t.Errorf("third event response_type=%v, want answer", third["response_type"])
	}
	// Fourth line: complete with done=true.
	var fourth map[string]any
	if err := json.Unmarshal([]byte(lines[3]), &fourth); err != nil {
		t.Fatalf("line 4 not JSON: %v", err)
	}
	if fourth["done"] != true {
		t.Errorf("fourth event done=%v, want true", fourth["done"])
	}
}