upload_test.go

  1package server
  2
  3import (
  4	"bytes"
  5	"encoding/json"
  6	"io"
  7	"log/slog"
  8	"mime/multipart"
  9	"net/http"
 10	"net/http/httptest"
 11	"os"
 12	"path/filepath"
 13	"strings"
 14	"testing"
 15
 16	"shelley.exe.dev/claudetool"
 17	"shelley.exe.dev/claudetool/browse"
 18	"shelley.exe.dev/loop"
 19)
 20
 21func TestUploadEndpoint(t *testing.T) {
 22	database, cleanup := setupTestDB(t)
 23	defer cleanup()
 24
 25	predictableService := loop.NewPredictableService()
 26	llmManager := &testLLMManager{service: predictableService}
 27	logger := slog.Default()
 28	server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
 29
 30	// Create a multipart form with a file
 31	body := &bytes.Buffer{}
 32	writer := multipart.NewWriter(body)
 33
 34	// Create a test file
 35	part, err := writer.CreateFormFile("file", "test.png")
 36	if err != nil {
 37		t.Fatalf("failed to create form file: %v", err)
 38	}
 39
 40	// Write some fake PNG content (just the magic header bytes)
 41	pngData := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}
 42	if _, err := part.Write(pngData); err != nil {
 43		t.Fatalf("failed to write file content: %v", err)
 44	}
 45	writer.Close()
 46
 47	req := httptest.NewRequest("POST", "/api/upload", body)
 48	req.Header.Set("Content-Type", writer.FormDataContentType())
 49	w := httptest.NewRecorder()
 50
 51	server.handleUpload(w, req)
 52
 53	if w.Code != http.StatusOK {
 54		t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String())
 55	}
 56
 57	var response map[string]string
 58	if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
 59		t.Fatalf("failed to parse response: %v", err)
 60	}
 61
 62	path, ok := response["path"]
 63	if !ok {
 64		t.Fatal("response missing 'path' field")
 65	}
 66
 67	// Verify the path is in the screenshot directory
 68	if !strings.HasPrefix(path, browse.ScreenshotDir) {
 69		t.Errorf("expected path to start with %s, got %s", browse.ScreenshotDir, path)
 70	}
 71
 72	// Verify the file has the correct extension
 73	if !strings.HasSuffix(path, ".png") {
 74		t.Errorf("expected path to end with .png, got %s", path)
 75	}
 76
 77	// Verify the file exists and contains our data
 78	data, err := os.ReadFile(path)
 79	if err != nil {
 80		t.Fatalf("failed to read uploaded file: %v", err)
 81	}
 82
 83	if !bytes.Equal(data, pngData) {
 84		t.Errorf("uploaded file content mismatch")
 85	}
 86
 87	// Clean up uploaded file
 88	os.Remove(path)
 89}
 90
 91func TestUploadEndpointMethodNotAllowed(t *testing.T) {
 92	database, cleanup := setupTestDB(t)
 93	defer cleanup()
 94
 95	predictableService := loop.NewPredictableService()
 96	llmManager := &testLLMManager{service: predictableService}
 97	logger := slog.Default()
 98	server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
 99
100	req := httptest.NewRequest("GET", "/api/upload", nil)
101	w := httptest.NewRecorder()
102
103	server.handleUpload(w, req)
104
105	if w.Code != http.StatusMethodNotAllowed {
106		t.Fatalf("expected status 405, got %d", w.Code)
107	}
108}
109
110func TestUploadEndpointNoFile(t *testing.T) {
111	database, cleanup := setupTestDB(t)
112	defer cleanup()
113
114	predictableService := loop.NewPredictableService()
115	llmManager := &testLLMManager{service: predictableService}
116	logger := slog.Default()
117	server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
118
119	// Create an empty multipart form
120	body := &bytes.Buffer{}
121	writer := multipart.NewWriter(body)
122	writer.Close()
123
124	req := httptest.NewRequest("POST", "/api/upload", body)
125	req.Header.Set("Content-Type", writer.FormDataContentType())
126	w := httptest.NewRecorder()
127
128	server.handleUpload(w, req)
129
130	if w.Code != http.StatusBadRequest {
131		t.Fatalf("expected status 400, got %d: %s", w.Code, w.Body.String())
132	}
133}
134
135func TestUploadedFileCanBeReadViaReadEndpoint(t *testing.T) {
136	database, cleanup := setupTestDB(t)
137	defer cleanup()
138
139	predictableService := loop.NewPredictableService()
140	llmManager := &testLLMManager{service: predictableService}
141	logger := slog.Default()
142	server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
143
144	// First, upload a file
145	body := &bytes.Buffer{}
146	writer := multipart.NewWriter(body)
147
148	part, err := writer.CreateFormFile("file", "test.jpg")
149	if err != nil {
150		t.Fatalf("failed to create form file: %v", err)
151	}
152
153	// Write some fake JPEG content
154	jpgData := []byte{0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 0x4A, 0x46, 0x49, 0x46}
155	if _, err := part.Write(jpgData); err != nil {
156		t.Fatalf("failed to write file content: %v", err)
157	}
158	writer.Close()
159
160	uploadReq := httptest.NewRequest("POST", "/api/upload", body)
161	uploadReq.Header.Set("Content-Type", writer.FormDataContentType())
162	uploadW := httptest.NewRecorder()
163
164	server.handleUpload(uploadW, uploadReq)
165
166	if uploadW.Code != http.StatusOK {
167		t.Fatalf("upload failed: %s", uploadW.Body.String())
168	}
169
170	var uploadResponse map[string]string
171	if err := json.Unmarshal(uploadW.Body.Bytes(), &uploadResponse); err != nil {
172		t.Fatalf("failed to parse upload response: %v", err)
173	}
174
175	path := uploadResponse["path"]
176
177	// Now try to read the file via the read endpoint
178	readReq := httptest.NewRequest("GET", "/api/read?path="+path, nil)
179	readW := httptest.NewRecorder()
180
181	server.handleRead(readW, readReq)
182
183	if readW.Code != http.StatusOK {
184		t.Fatalf("read failed with status %d: %s", readW.Code, readW.Body.String())
185	}
186
187	// Verify content type
188	contentType := readW.Header().Get("Content-Type")
189	if contentType != "image/jpeg" {
190		t.Errorf("expected Content-Type image/jpeg, got %s", contentType)
191	}
192
193	// Verify content
194	readData, err := io.ReadAll(readW.Body)
195	if err != nil {
196		t.Fatalf("failed to read response body: %v", err)
197	}
198
199	if !bytes.Equal(readData, jpgData) {
200		t.Errorf("read content mismatch")
201	}
202
203	// Clean up
204	os.Remove(path)
205}
206
207func TestUploadPreservesFileExtension(t *testing.T) {
208	database, cleanup := setupTestDB(t)
209	defer cleanup()
210
211	predictableService := loop.NewPredictableService()
212	llmManager := &testLLMManager{service: predictableService}
213	logger := slog.Default()
214	server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
215
216	testCases := []struct {
217		filename string
218		wantExt  string
219	}{
220		{"photo.png", ".png"},
221		{"image.jpeg", ".jpeg"},
222		{"screenshot.gif", ".gif"},
223		{"document.pdf", ".pdf"},
224		{"noextension", ""},
225	}
226
227	for _, tc := range testCases {
228		t.Run(tc.filename, func(t *testing.T) {
229			body := &bytes.Buffer{}
230			writer := multipart.NewWriter(body)
231
232			part, err := writer.CreateFormFile("file", tc.filename)
233			if err != nil {
234				t.Fatalf("failed to create form file: %v", err)
235			}
236			part.Write([]byte("test content"))
237			writer.Close()
238
239			req := httptest.NewRequest("POST", "/api/upload", body)
240			req.Header.Set("Content-Type", writer.FormDataContentType())
241			w := httptest.NewRecorder()
242
243			server.handleUpload(w, req)
244
245			if w.Code != http.StatusOK {
246				t.Fatalf("expected status 200, got %d", w.Code)
247			}
248
249			var response map[string]string
250			if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
251				t.Fatalf("failed to parse response: %v", err)
252			}
253
254			path := response["path"]
255			ext := filepath.Ext(path)
256			if ext != tc.wantExt {
257				t.Errorf("expected extension %q, got %q", tc.wantExt, ext)
258			}
259
260			// Clean up
261			os.Remove(path)
262		})
263	}
264}