44 lines
1.4 KiB
Go
44 lines
1.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())
|
|
}
|
|
}
|