digital-psychology/apps/api/internal/llm/deepseek/client_test.go

package deepseek

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/yuxingu/digital-psychology/apps/api/internal/config"
)

func TestChat_success(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer test-key" {
			t.Fatalf("auth header: %s", r.Header.Get("Authorization"))
		}
		var req chatRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		if req.Model != "deepseek-chat" {
			t.Fatalf("model=%s", req.Model)
		}
		_ = json.NewEncoder(w).Encode(chatResponse{
			Choices: []struct {
				Message Message `json:"message"`
			}{{Message: Message{Role: "assistant", Content: "你好，我是成长助手"}}},
		})
	}))
	defer srv.Close()

	c := New(config.DeepSeekConfig{
		APIKey: "test-key", BaseURL: srv.URL, Model: "deepseek-chat", TimeoutSec: 5,
	})
	out, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}})
	if err != nil {
		t.Fatal(err)
	}
	if out != "你好，我是成长助手" {
		t.Fatalf("got %q", out)
	}
}

func TestChatStream_success(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req chatRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		if !req.Stream {
			t.Fatal("expected stream=true")
		}
		w.Header().Set("Content-Type", "text/event-stream")
		flusher := w.(http.Flusher)
		_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"你好\"}}]}\n\n"))
		flusher.Flush()
		_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"世界\"}}]}\n\n"))
		flusher.Flush()
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
		flusher.Flush()
	}))
	defer srv.Close()

	c := New(config.DeepSeekConfig{
		APIKey: "test-key", BaseURL: srv.URL, Model: "deepseek-chat", TimeoutSec: 5,
	})
	var got string
	out, err := c.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, func(d string) error {
		got += d
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if out != "你好世界" || got != "你好世界" {
		t.Fatalf("out=%q got=%q", out, got)
	}
}