61 lines
2.4 KiB
Go
61 lines
2.4 KiB
Go
package mockprovider
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestChatSuccessAndErrorMatrix(t *testing.T) {
|
|
handler := Handler{}
|
|
for _, test := range []struct {
|
|
prompt string
|
|
status int
|
|
}{{"hello", 200}, {"mock:429", 429}, {"mock:500", 500}, {"mock:400", 400}, {"mock:401", 401}, {"mock:policy", 400}} {
|
|
t.Run(test.prompt, func(t *testing.T) {
|
|
body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"content": test.prompt}}})
|
|
request := httptest.NewRequest(http.MethodPost, "http://mock/v1/chat/completions", bytes.NewReader(body))
|
|
response := httptest.NewRecorder()
|
|
handler.ServeHTTP(response, request)
|
|
if response.Code != test.status {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
func TestImageSuccess(t *testing.T) {
|
|
var body bytes.Buffer
|
|
writer := multipart.NewWriter(&body)
|
|
_ = writer.WriteField("prompt", "edit")
|
|
part, _ := writer.CreateFormFile("image", "input.png")
|
|
part.Write(mockPNG())
|
|
writer.Close()
|
|
request := httptest.NewRequest(http.MethodPost, "http://mock/v1/images/edits", &body)
|
|
request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
response := httptest.NewRecorder()
|
|
Handler{}.ServeHTTP(response, request)
|
|
if response.Code != 200 || !bytes.Contains(response.Body.Bytes(), []byte("b64_json")) {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestImagesGenerationAndGeminiEndpoints(t *testing.T) {
|
|
imageRequest := httptest.NewRequest(http.MethodPost, "http://mock/v1/images/generations", bytes.NewBufferString(`{"model":"mock","prompt":"draw","n":1}`))
|
|
imageRequest.Header.Set("Content-Type", "application/json")
|
|
imageResponse := httptest.NewRecorder()
|
|
Handler{}.ServeHTTP(imageResponse, imageRequest)
|
|
if imageResponse.Code != http.StatusOK || !bytes.Contains(imageResponse.Body.Bytes(), []byte("b64_json")) {
|
|
t.Fatalf("images status=%d body=%s", imageResponse.Code, imageResponse.Body.String())
|
|
}
|
|
|
|
geminiRequest := httptest.NewRequest(http.MethodPost, "http://mock/v1/models/mock:generateContent", bytes.NewBufferString(`{"contents":[{"parts":[{"text":"hello"}]}]}`))
|
|
geminiResponse := httptest.NewRecorder()
|
|
Handler{}.ServeHTTP(geminiResponse, geminiRequest)
|
|
if geminiResponse.Code != http.StatusOK || !bytes.Contains(geminiResponse.Body.Bytes(), []byte("candidates")) {
|
|
t.Fatalf("gemini status=%d body=%s", geminiResponse.Code, geminiResponse.Body.String())
|
|
}
|
|
}
|