Some checks failed
Build & Sign Wraith / Build Windows + Sign (push) Has been cancelled
- OAuth PKCE flow for Max subscription auth (no API key needed) - Claude API client with SSE streaming (Messages API v1) - 16 tool definitions: terminal, SFTP, RDP, session management - Tool dispatch router mapping to existing Wraith services - Conversation manager with SQLite persistence - Terminal output ring buffer for AI context - RDP screenshot encoder (RGBA → JPEG with downscaling) - Wired into Wails app as AIService Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
124 lines
3.2 KiB
Go
124 lines
3.2 KiB
Go
package ai
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"testing"
|
|
)
|
|
|
|
func TestBuildRequestBody(t *testing.T) {
|
|
messages := []Message{
|
|
{
|
|
Role: "user",
|
|
Content: []ContentBlock{
|
|
{Type: "text", Text: "Hello"},
|
|
},
|
|
},
|
|
}
|
|
tools := []Tool{
|
|
{
|
|
Name: "test_tool",
|
|
Description: "A test tool",
|
|
InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
|
|
},
|
|
}
|
|
|
|
body, err := BuildRequestBody(messages, tools, "You are helpful.", "")
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
var parsed map[string]interface{}
|
|
if err := json.Unmarshal(body, &parsed); err != nil {
|
|
t.Fatalf("invalid JSON: %v", err)
|
|
}
|
|
|
|
// Check model
|
|
if m, ok := parsed["model"].(string); !ok || m != defaultModel {
|
|
t.Errorf("expected model %s, got %v", defaultModel, parsed["model"])
|
|
}
|
|
|
|
// Check stream
|
|
if s, ok := parsed["stream"].(bool); !ok || !s {
|
|
t.Errorf("expected stream true, got %v", parsed["stream"])
|
|
}
|
|
|
|
// Check system prompt
|
|
if s, ok := parsed["system"].(string); !ok || s != "You are helpful." {
|
|
t.Errorf("expected system prompt, got %v", parsed["system"])
|
|
}
|
|
|
|
// Check max_tokens
|
|
if mt, ok := parsed["max_tokens"].(float64); !ok || int(mt) != 8192 {
|
|
t.Errorf("expected max_tokens 8192, got %v", parsed["max_tokens"])
|
|
}
|
|
|
|
// Check messages array exists
|
|
msgs, ok := parsed["messages"].([]interface{})
|
|
if !ok || len(msgs) != 1 {
|
|
t.Errorf("expected 1 message, got %v", parsed["messages"])
|
|
}
|
|
|
|
// Check tools array exists
|
|
tls, ok := parsed["tools"].([]interface{})
|
|
if !ok || len(tls) != 1 {
|
|
t.Errorf("expected 1 tool, got %v", parsed["tools"])
|
|
}
|
|
}
|
|
|
|
func TestParseSSELine(t *testing.T) {
|
|
tests := []struct {
|
|
input string
|
|
wantType string
|
|
wantData string
|
|
}{
|
|
{"event: content_block_delta", "event", "content_block_delta"},
|
|
{"data: {\"delta\":{\"text\":\"hello\"}}", "data", "{\"delta\":{\"text\":\"hello\"}}"},
|
|
{"", "", ""},
|
|
{"random line", "", ""},
|
|
{"event: message_stop", "event", "message_stop"},
|
|
{"data: [DONE]", "data", "[DONE]"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.input, func(t *testing.T) {
|
|
gotType, gotData := ParseSSELine(tt.input)
|
|
if gotType != tt.wantType {
|
|
t.Errorf("ParseSSELine(%q) type = %q, want %q", tt.input, gotType, tt.wantType)
|
|
}
|
|
if gotData != tt.wantData {
|
|
t.Errorf("ParseSSELine(%q) data = %q, want %q", tt.input, gotData, tt.wantData)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAuthHeader(t *testing.T) {
|
|
// Test with API key only (no OAuth)
|
|
client := NewClaudeClient(nil, "")
|
|
client.SetAPIKey("sk-test-12345")
|
|
|
|
req, _ := http.NewRequest("POST", "https://example.com", nil)
|
|
if err := client.setAuthHeader(req); err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
if got := req.Header.Get("x-api-key"); got != "sk-test-12345" {
|
|
t.Errorf("expected x-api-key sk-test-12345, got %q", got)
|
|
}
|
|
if got := req.Header.Get("Authorization"); got != "" {
|
|
t.Errorf("expected no Authorization header, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestAuthHeaderNoAuth(t *testing.T) {
|
|
// No OAuth, no API key — should return error
|
|
client := NewClaudeClient(nil, "")
|
|
|
|
req, _ := http.NewRequest("POST", "https://example.com", nil)
|
|
err := client.setAuthHeader(req)
|
|
if err == nil {
|
|
t.Error("expected error when no auth method available")
|
|
}
|
|
}
|