go-deepseek/completion_test.go

package deepseek

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"reflect"
	"strings"
	"testing"
)

func TestCompletionReqBuilder(t *testing.T) {
	req := NewCompletionReq("chat-id").
		SetParentMessageID(42).
		SetModelType(ModelTypeExpert).
		AddPrompt("hello").
		AddPrompt(" world").
		BuildPrompt().
		SetThinkingEnabled(true).
		SetSearchEnabled(true)
	req.RefFileIDs = []string{"file-id"}

	if req.ChatSessionID != "chat-id" || req.ModelType != ModelTypeExpert || req.Prompt != "hello world" {
		t.Fatalf("request = %#v", req)
	}
	if req.ParentMessageID == nil || *req.ParentMessageID != 42 {
		t.Fatalf("parent message ID = %v", req.ParentMessageID)
	}
	if !req.ThinkingEnabled || !req.SearchEnabled || !reflect.DeepEqual(req.RefFileIDs, []string{"file-id"}) {
		t.Fatalf("request options = %#v", req)
	}

	req.AddPrompt("new").BuildPrompt()
	if req.Prompt != "new" {
		t.Fatalf("rebuilt prompt = %q, want new", req.Prompt)
	}
}

func TestCompletionWithContextPreservesRequestOptions(t *testing.T) {
	const challenge = "2f90572ad390d758b5e55b3bb74f14722166388023b3b28876d056a358591197"
	const salt = "2eeb8f3a703002bfca70"
	const expectedAnswer = uint64(61830)

	client := testAPI(func(r *http.Request) (*http.Response, error) {
		switch r.URL.Path {
		case "/chat/create_pow_challenge":
			return testHTTPResponse(`{"code":0,"data":{"biz_code":0,"biz_data":{"challenge":{"algorithm":"DeepSeekHashV1","challenge":"` + challenge + `","salt":"` + salt + `","signature":"signature","difficulty":144000,"expire_at":1785483643587,"target_path":"/api/v0/chat/completion"}}}}`), nil
		case "/chat/completion":
			encoded := r.Header.Get("x-ds-pow-response")
			raw, err := base64.URLEncoding.DecodeString(encoded)
			if err != nil {
				t.Errorf("decode PoW header: %v", err)
			}
			var proof struct {
				Answer     uint64 `json:"answer"`
				Signature  string `json:"signature"`
				TargetPath string `json:"target_path"`
			}
			if err := json.Unmarshal(raw, &proof); err != nil {
				t.Errorf("decode PoW JSON: %v", err)
			}
			if proof.Answer != expectedAnswer || proof.Signature != "signature" || proof.TargetPath != "/api/v0/chat/completion" {
				t.Errorf("proof = %#v", proof)
			}

			var body CompletionReq
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				t.Errorf("decode completion: %v", err)
			}
			if body.ChatSessionID != "chat-id" || body.ParentMessageID == nil || *body.ParentMessageID != 9 || body.ModelType != ModelTypeExpert || body.Prompt != "prompt" || !body.ThinkingEnabled || !body.SearchEnabled || !reflect.DeepEqual(body.RefFileIDs, []string{"file-id"}) {
				t.Errorf("completion body = %#v", body)
			}
			resp := testHTTPResponse("event: close\ndata: {}\n\n")
			resp.Header.Set("Content-Type", "text/event-stream")
			return resp, nil
		default:
			t.Fatalf("unexpected path %q", r.URL.Path)
			return nil, nil
		}
	})

	parentID := uint64(9)
	body := CompletionReq{
		ChatSessionID:   "chat-id",
		ParentMessageID: &parentID,
		ModelType:       ModelTypeExpert,
		Prompt:          "prompt",
		RefFileIDs:      []string{"file-id"},
		ThinkingEnabled: true,
		SearchEnabled:   true,
	}
	resp, err := client.CompletionWithContext(context.Background(), body)
	if err != nil {
		t.Fatalf("CompletionWithContext() error = %v", err)
	}
	defer resp.Body.Close()
}

type trackingReadCloser struct {
	io.Reader
	closed bool
}

func (r *trackingReadCloser) Close() error {
	r.closed = true
	return nil
}

func TestReadStream(t *testing.T) {
	stream := strings.Join([]string{
		`event: ready`,
		`data: {"request_message_id":1,"response_message_id":2,"model_type":"default"}`,
		``,
		`data: {"p":"response/content","o":"APPEND","v":"hel"}`,
		``,
		`data: {"v":"lo"}`,
		``,
		`event: close`,
		`data: {}`,
		``,
	}, "\n")
	body := &trackingReadCloser{Reader: strings.NewReader(stream)}
	resp := &http.Response{Body: body}
	var chunks []string

	state, err := ReadStream(context.Background(), resp, func(s string) {
		chunks = append(chunks, s)
	})
	if err != nil {
		t.Fatalf("ReadStream() error = %v", err)
	}
	if !body.closed {
		t.Fatal("response body was not closed")
	}
	if got := state.Content.String(); got != "hello" {
		t.Fatalf("content = %q, want hello", got)
	}
	if !reflect.DeepEqual(chunks, []string{"hel", "lo"}) {
		t.Fatalf("chunks = %#v", chunks)
	}
	if state.RequestMessageID != 1 || state.ResponseMessageID != 2 || !state.Closed {
		t.Fatalf("state = %#v", state)
	}
}

func TestReadStreamAsStringReturnsFinalSetValue(t *testing.T) {
	stream := strings.Join([]string{
		`data: {"p":"response/content","o":"APPEND","v":"obsolete value"}`,
		``,
		`data: {"p":"response/content","o":"SET","v":"final"}`,
		``,
		`event: close`,
		`data: {}`,
		``,
	}, "\n")
	body := &trackingReadCloser{Reader: strings.NewReader(stream)}

	got, err := ReadStreamAsString(context.Background(), &http.Response{Body: body})
	if err != nil {
		t.Fatalf("ReadStreamAsString() error = %v", err)
	}
	if got != "final" {
		t.Fatalf("ReadStreamAsString() = %q, want final", got)
	}
	if !body.closed {
		t.Fatal("response body was not closed")
	}
}