wraith/internal/ai/screenshot_test.go
Vantz Stockwell 7ee5321d69
Some checks failed
Build & Sign Wraith / Build Windows + Sign (push) Has been cancelled
feat: AI copilot backend — OAuth PKCE, Claude API streaming, 16 tools, conversations
- 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>
2026-03-17 09:09:23 -04:00

119 lines
3.3 KiB
Go

package ai
import (
"bytes"
"image/jpeg"
"testing"
)
func TestEncodeScreenshot(t *testing.T) {
width, height := 100, 80
// Create a test RGBA buffer (red pixels)
rgba := make([]byte, width*height*4)
for i := 0; i < len(rgba); i += 4 {
rgba[i+0] = 255 // R
rgba[i+1] = 0 // G
rgba[i+2] = 0 // B
rgba[i+3] = 255 // A
}
jpegData, err := EncodeScreenshot(rgba, width, height, 200, 200, 80)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Check JPEG magic bytes (FF D8)
if len(jpegData) < 2 {
t.Fatal("JPEG data too small")
}
if jpegData[0] != 0xFF || jpegData[1] != 0xD8 {
t.Errorf("expected JPEG magic bytes FF D8, got %02X %02X", jpegData[0], jpegData[1])
}
// Decode and verify dimensions (no downscale needed in this case)
img, err := jpeg.Decode(bytes.NewReader(jpegData))
if err != nil {
t.Fatalf("decode JPEG: %v", err)
}
bounds := img.Bounds()
if bounds.Dx() != width || bounds.Dy() != height {
t.Errorf("expected %dx%d, got %dx%d", width, height, bounds.Dx(), bounds.Dy())
}
}
func TestEncodeScreenshotDownscale(t *testing.T) {
srcWidth, srcHeight := 1920, 1080
maxWidth, maxHeight := 1280, 720
// Create a test RGBA buffer (gradient)
rgba := make([]byte, srcWidth*srcHeight*4)
for y := 0; y < srcHeight; y++ {
for x := 0; x < srcWidth; x++ {
idx := (y*srcWidth + x) * 4
rgba[idx+0] = byte(x % 256) // R
rgba[idx+1] = byte(y % 256) // G
rgba[idx+2] = byte((x + y) % 256) // B
rgba[idx+3] = 255 // A
}
}
jpegData, err := EncodeScreenshot(rgba, srcWidth, srcHeight, maxWidth, maxHeight, 75)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Check JPEG magic bytes
if jpegData[0] != 0xFF || jpegData[1] != 0xD8 {
t.Errorf("expected JPEG magic bytes FF D8, got %02X %02X", jpegData[0], jpegData[1])
}
// Decode and verify dimensions are within max bounds
img, err := jpeg.Decode(bytes.NewReader(jpegData))
if err != nil {
t.Fatalf("decode JPEG: %v", err)
}
bounds := img.Bounds()
if bounds.Dx() > maxWidth {
t.Errorf("output width %d exceeds max %d", bounds.Dx(), maxWidth)
}
if bounds.Dy() > maxHeight {
t.Errorf("output height %d exceeds max %d", bounds.Dy(), maxHeight)
}
// Should maintain 16:9 ratio approximately
expectedWidth := 1280
expectedHeight := 720
if bounds.Dx() != expectedWidth || bounds.Dy() != expectedHeight {
t.Errorf("expected %dx%d, got %dx%d", expectedWidth, expectedHeight, bounds.Dx(), bounds.Dy())
}
}
func TestEncodeScreenshotBufferTooSmall(t *testing.T) {
_, err := EncodeScreenshot([]byte{0, 0, 0, 0}, 100, 100, 200, 200, 75)
if err == nil {
t.Error("expected error for buffer too small")
}
}
func TestFitDimensions(t *testing.T) {
tests := []struct {
srcW, srcH, maxW, maxH int
wantW, wantH int
}{
{1920, 1080, 1280, 720, 1280, 720}, // 16:9 fits exactly
{1920, 1080, 800, 600, 800, 450}, // width-constrained
{1080, 1920, 800, 600, 337, 600}, // height-constrained (portrait)
{100, 100, 200, 200, 200, 200}, // smaller than max (but func called with > check)
}
for _, tt := range tests {
w, h := fitDimensions(tt.srcW, tt.srcH, tt.maxW, tt.maxH)
if w != tt.wantW || h != tt.wantH {
t.Errorf("fitDimensions(%d,%d,%d,%d) = %d,%d, want %d,%d",
tt.srcW, tt.srcH, tt.maxW, tt.maxH, w, h, tt.wantW, tt.wantH)
}
}
}