initial commit
This commit is contained in:
+194
@@ -0,0 +1,194 @@
|
||||
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 TestCreateChatWithContext(t *testing.T) {
|
||||
client := testAPI(func(r *http.Request) (*http.Response, error) {
|
||||
if r.URL.Path != "/chat_session/create" {
|
||||
t.Errorf("path = %q", r.URL.Path)
|
||||
}
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Errorf("read body: %v", err)
|
||||
}
|
||||
if string(body) != `{}` {
|
||||
t.Errorf("body = %s, want {}", body)
|
||||
}
|
||||
return testHTTPResponse(`{"code":0,"data":{"biz_code":0,"biz_data":{"chat_session":{"id":"chat-id","seq_id":7,"model_type":"default"},"ttl_seconds":60}}}`), nil
|
||||
})
|
||||
|
||||
got, err := client.CreateChatWithContext(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("CreateChatWithContext() error = %v", err)
|
||||
}
|
||||
if got.ID != "chat-id" || got.SeqID != 7 || got.ModelType != "default" {
|
||||
t.Fatalf("session = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
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 := NewClient().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 := NewClient().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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user