Repository: ollama/ollama
File: openai/openai_test.go (Go)

Code before:
const prefix = `data:image/jpeg;base64,`
const image = `iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=`
const imageURL = prefix + image

func TestMiddlewareRequests(t *testing.T) {
	type testCase struct {
		Name     string
		Method   string
		Path     string
		Handler  func() gin.HandlerFunc
		Setup    func(t *testing.T, req *http.Request)
		Expected func(t *testing.T, req *http.Request)
	}

	var capturedRequest *http.Request

	captureRequestMiddleware := func() gin.HandlerFunc {
		return func(c *gin.Context) {
			bodyBytes, _ := io.ReadAll(c.Request.Body)
			c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			capturedRequest = c.Request
			c.Next()
		}
	}

	testCases := []testCase{
		{
			Name:    "chat handler",
			Method:  http.MethodPost,
			Path:    "/api/chat",
			Handler: ChatMiddleware,
			Setup: func(t *testing.T, req *http.Request) {
				body := ChatCompletionRequest{
					Model:    "test-model",
					Messages: []Message{{Role: "user", Content: "Hello"}},
				}

				bodyBytes, _ := json.Marshal(body)

				req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
				req.Header.Set("Content-Type", "application/json")
			},
			Expected: func(t *testing.T, req *http.Request) {
				var chatReq api.ChatRequest
				if err := json.NewDecoder(req.Body).Decode(&chatReq); err != nil {
					t.Fatal(err)
				}

				if chatReq.Messages[0].Role != "user" {
					t.Fatalf("expected 'user', got %s", chatReq.Messages[0].Role)
				}

Review comment (refactor):
As a follow up, we could simplify this to:

```
type testCase struct {
  name string // name of case
  input string // incoming body
  request api.ChatRequest // proxied chat request
  err error
}
```

Code after:
const prefix = `data:image/jpeg;base64,`
const image = `iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=`
const imageURL = prefix + image

func prepareRequest(req *http.Request, body any) {
	bodyBytes, _ := json.Marshal(body)
	req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
	req.Header.Set("Content-Type", "application/json")
}

func captureRequestMiddleware(capturedRequest any) gin.HandlerFunc {
	return func(c *gin.Context) {
		bodyBytes, _ := io.ReadAll(c.Request.Body)
		c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
		err := json.Unmarshal(bodyBytes, capturedRequest)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, "failed to unmarshal request")
		}
		c.Next()
	}
}

func TestChatMiddleware(t *testing.T) {
	type testCase struct {
		Name     string
		Setup    func(t *testing.T, req *http.Request)
		Expected func(t *testing.T, req *api.ChatRequest, resp *httptest.ResponseRecorder)
	}

	var capturedRequest *api.ChatRequest

	testCases := []testCase{
		{
			Name: "chat handler",
			Setup: func(t *testing.T, req *http.Request) {
				body := ChatCompletionRequest{
					Model:    "test-model",
					Messages: []Message{{Role: "user", Content: "Hello"}},
				}
				prepareRequest(req, body)
			},
			Expected: func(t *testing.T, req *api.ChatRequest, resp *httptest.ResponseRecorder) {
				if resp.Code != http.StatusOK {
					t.Fatalf("expected 200, got %d", resp.Code)
				}

				if req.Messages[0].Role != "user" {
					t.Fatalf("expected 'user', got %s", req.Messages[0].Role)
				}

				if req.Messages[0].Content != "Hello" {