LaodingBot/internal/transport/webui/bot_test.go

package webui

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"laodingbot/internal/config"
	"laodingbot/internal/llm"
)

func newTestBot(t *testing.T, maxUploadBytes int64) *Bot {
	t.Helper()
	b, err := NewBot(config.WebUIConfig{ListenAddr: ":8090", MaxUploadBytes: maxUploadBytes}, nil)
	if err != nil {
		t.Fatalf("NewBot failed: %v", err)
	}
	return b
}

func TestHandleChatSuccess(t *testing.T) {
	b := newTestBot(t, 1024*1024)
	b.chatHandler = func(_ context.Context, msg IncomingMessage) (string, error) {
		if msg.ChatID != "s1" || msg.UserID != "u1" || msg.Text != "hello" {
			t.Fatalf("unexpected message: %+v", msg)
		}
		return "ok", nil
	}

	body := strings.NewReader(`{"text":"hello","session_id":"s1","user_id":"u1"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/chat", body)
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	b.handleChat(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String())
	}

	var out chatResponse
	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
		t.Fatalf("decode response failed: %v", err)
	}
	if out.Reply != "ok" || out.SessionID != "s1" || out.UserID != "u1" {
		t.Fatalf("unexpected response: %+v", out)
	}
}

func TestHandleChatMissingText(t *testing.T) {
	b := newTestBot(t, 1024*1024)
	b.chatHandler = func(_ context.Context, _ IncomingMessage) (string, error) { return "", nil }

	body := strings.NewReader(`{"text":"   "}`)
	req := httptest.NewRequest(http.MethodPost, "/api/chat", body)
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	b.handleChat(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", w.Code)
	}
}

func TestHandleUploadSuccess(t *testing.T) {
	b := newTestBot(t, 1024*1024)
	b.uploadHandler = func(_ context.Context, chatID, userID string, files []llm.InputFile) ([]string, error) {
		if chatID != "s1" || userID != "u1" {
			t.Fatalf("unexpected ids chat=%s user=%s", chatID, userID)
		}
		if len(files) != 1 {
			t.Fatalf("unexpected files len=%d", len(files))
		}
		if files[0].FileName != "doc.pdf" || len(files[0].Content) == 0 {
			t.Fatalf("unexpected file payload: %+v", files[0])
		}
		return []string{"file_123"}, nil
	}

	var payload bytes.Buffer
	writer := multipart.NewWriter(&payload)
	_ = writer.WriteField("session_id", "s1")
	_ = writer.WriteField("user_id", "u1")
	fw, err := writer.CreateFormFile("file", "doc.pdf")
	if err != nil {
		t.Fatalf("CreateFormFile failed: %v", err)
	}
	_, _ = fw.Write([]byte("pdf-content"))
	_ = writer.Close()

	req := httptest.NewRequest(http.MethodPost, "/api/upload", &payload)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	w := httptest.NewRecorder()

	b.handleUpload(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String())
	}

	var out uploadResponse
	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
		t.Fatalf("decode response failed: %v", err)
	}
	if out.FileID != "file_123" || out.SessionID != "s1" || out.UserID != "u1" {
		t.Fatalf("unexpected response: %+v", out)
	}
}

func TestHandleUploadTooLarge(t *testing.T) {
	b := newTestBot(t, 3)
	b.uploadHandler = func(_ context.Context, _ string, _ string, _ []llm.InputFile) ([]string, error) {
		return []string{"file_should_not_reach"}, nil
	}

	var payload bytes.Buffer
	writer := multipart.NewWriter(&payload)
	fw, err := writer.CreateFormFile("file", "a.txt")
	if err != nil {
		t.Fatalf("CreateFormFile failed: %v", err)
	}
	_, _ = fw.Write([]byte("12345"))
	_ = writer.Close()

	req := httptest.NewRequest(http.MethodPost, "/api/upload", &payload)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	w := httptest.NewRecorder()

	b.handleUpload(w, req)
	if w.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("expected 413, got %d body=%s", w.Code, w.Body.String())
	}
}

func TestHandleUploadMissingFile(t *testing.T) {
	b := newTestBot(t, 1024*1024)
	b.uploadHandler = func(_ context.Context, _ string, _ string, _ []llm.InputFile) ([]string, error) {
		return []string{"file_should_not_reach"}, nil
	}

	var payload bytes.Buffer
	writer := multipart.NewWriter(&payload)
	_ = writer.WriteField("session_id", "s1")
	_ = writer.Close()

	req := httptest.NewRequest(http.MethodPost, "/api/upload", &payload)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	w := httptest.NewRecorder()

	b.handleUpload(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", w.Code)
	}
}

func TestHandleChatStreamSuccess(t *testing.T) {
	b := newTestBot(t, 1024*1024)
	b.streamChatHandler = func(_ context.Context, msg IncomingMessage, cb StreamEventCallback) (string, error) {
		if msg.ChatID != "s1" || msg.UserID != "u1" || msg.Text != "hello" {
			t.Fatalf("unexpected message: %+v", msg)
		}
		if err := cb(StreamEvent{Type: StreamEventTypeThought, Content: "thinking", Step: 1}); err != nil {
			return "", err
		}
		if err := cb(StreamEvent{Type: StreamEventTypeToolCall, Content: "{\"input\":\"pwd\"}", Step: 1, ToolName: "shell"}); err != nil {
			return "", err
		}
		if err := cb(StreamEvent{Type: StreamEventTypeToolResult, Content: "C:/Project", Step: 1, ToolName: "shell"}); err != nil {
			return "", err
		}
		if err := cb(StreamEvent{Type: StreamEventTypeFinal, Content: "done", Step: 2}); err != nil {
			return "", err
		}
		return "done", nil
	}

	body := strings.NewReader(`{"text":"hello","session_id":"s1","user_id":"u1"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/chat/stream", body)
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	b.handleChatStream(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String())
	}
	if got := w.Header().Get("Content-Type"); got != "text/event-stream" {
		t.Fatalf("expected text/event-stream, got %q", got)
	}

	var events []StreamEvent
	chunks := strings.Split(strings.TrimSpace(w.Body.String()), "\n\n")
	for _, chunk := range chunks {
		line := strings.TrimSpace(chunk)
		if line == "" {
			continue
		}
		if !strings.HasPrefix(line, "data: ") {
			t.Fatalf("invalid sse line: %q", line)
		}
		payload := strings.TrimPrefix(line, "data: ")
		var ev StreamEvent
		if err := json.Unmarshal([]byte(payload), &ev); err != nil {
			t.Fatalf("unmarshal stream event failed: %v payload=%s", err, payload)
		}
		events = append(events, ev)
	}

	if len(events) != 4 {
		t.Fatalf("expected 4 events, got %d", len(events))
	}
	if events[0].Type != StreamEventTypeThought {
		t.Fatalf("event[0] type mismatch: %s", events[0].Type)
	}
	if events[1].Type != StreamEventTypeToolCall || events[1].ToolName != "shell" {
		t.Fatalf("event[1] mismatch: %+v", events[1])
	}
	if events[2].Type != StreamEventTypeToolResult || events[2].ToolName != "shell" {
		t.Fatalf("event[2] mismatch: %+v", events[2])
	}
	if events[3].Type != StreamEventTypeFinal || events[3].Content != "done" {
		t.Fatalf("event[3] mismatch: %+v", events[3])
	}
}

func TestHandleChatStreamHandlerError(t *testing.T) {
	b := newTestBot(t, 1024*1024)
	b.streamChatHandler = func(_ context.Context, _ IncomingMessage, _ StreamEventCallback) (string, error) {
		return "", errors.New("boom")
	}

	body := strings.NewReader(`{"text":"hello"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/chat/stream", body)
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	b.handleChatStream(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	respBody := w.Body.String()
	if !strings.Contains(respBody, `"type":"error"`) {
		t.Fatalf("expected error event in stream, body=%q", respBody)
	}
	if !strings.Contains(respBody, "stream error: boom") {
		t.Fatalf("expected error detail in stream, body=%q", respBody)
	}
}

func TestHandleChatStreamValidation(t *testing.T) {
	t.Run("method not allowed", func(t *testing.T) {
		b := newTestBot(t, 1024*1024)
		req := httptest.NewRequest(http.MethodGet, "/api/chat/stream", nil)
		w := httptest.NewRecorder()

		b.handleChatStream(w, req)
		if w.Code != http.StatusMethodNotAllowed {
			t.Fatalf("expected 405, got %d", w.Code)
		}
	})

	t.Run("content type must be json", func(t *testing.T) {
		b := newTestBot(t, 1024*1024)
		req := httptest.NewRequest(http.MethodPost, "/api/chat/stream", strings.NewReader(`{"text":"hello"}`))
		w := httptest.NewRecorder()

		b.handleChatStream(w, req)
		if w.Code != http.StatusBadRequest {
			t.Fatalf("expected 400, got %d", w.Code)
		}
	})

	t.Run("handler not ready", func(t *testing.T) {
		b := newTestBot(t, 1024*1024)
		req := httptest.NewRequest(http.MethodPost, "/api/chat/stream", strings.NewReader(`{"text":"hello"}`))
		req.Header.Set("Content-Type", "application/json")
		w := httptest.NewRecorder()

		b.handleChatStream(w, req)
		if w.Code != http.StatusInternalServerError {
			t.Fatalf("expected 500, got %d", w.Code)
		}
	})

	t.Run("text required", func(t *testing.T) {
		b := newTestBot(t, 1024*1024)
		b.streamChatHandler = func(_ context.Context, _ IncomingMessage, _ StreamEventCallback) (string, error) {
			return "", nil
		}
		req := httptest.NewRequest(http.MethodPost, "/api/chat/stream", strings.NewReader(`{"text":"   "}`))
		req.Header.Set("Content-Type", "application/json")
		w := httptest.NewRecorder()

		b.handleChatStream(w, req)
		if w.Code != http.StatusBadRequest {
			t.Fatalf("expected 400, got %d", w.Code)
		}
	})
}