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

Code before:

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

				if chatResp.Object != "chat.completion" {
					t.Fatalf("expected chat.completion, got %s", chatResp.Object)
				}

				if chatResp.Choices[0].Message.Content != "Hello!" {
					t.Fatalf("expected Hello!, got %s", chatResp.Choices[0].Message.Content)
				}
			},
		},
		{
			Name:     "completions handler",
			Method:   http.MethodPost,
			Path:     "/api/generate",
			TestPath: "/api/generate",
			Handler:  CompletionsMiddleware,
			Endpoint: func(c *gin.Context) {
				c.JSON(http.StatusOK, api.GenerateResponse{
					Response: "Hello!",
				})
			},
			Setup: func(t *testing.T, req *http.Request) {
				body := CompletionRequest{
					Model:  "test-model",
					Prompt: "Hello",
				}

				bodyBytes, _ := json.Marshal(body)

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

				assert.Equal(t, "text_completion", completionResp.Object)
				assert.Equal(t, "Hello!", completionResp.Choices[0].Text)
			},
		},

Review comment (suggestion):
We could probably use 2-3 more cases:

1. invalid types
2. params: temperature etc
3. forwarding an error

Code after:
				}

				bodyBytes, _ := json.Marshal(body)

				req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
				req.Header.Set("Content-Type", "application/json")
			},
			Expected: func(t *testing.T, resp *httptest.ResponseRecorder) {
				assert.Equal(t, http.StatusOK, resp.Code)

				var chatResp ChatCompletion
				if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
					t.Fatal(err)
				}

				if chatResp.Object != "chat.completion" {
					t.Fatalf("expected chat.completion, got %s", chatResp.Object)
				}

				if chatResp.Choices[0].Message.Content != "Hello!" {
					t.Fatalf("expected Hello!, got %s", chatResp.Choices[0].Message.Content)
				}
			},
		},
		{
			Name:     "completions handler",
			Method:   http.MethodPost,
			Path:     "/api/generate",
			TestPath: "/api/generate",
			Handler:  CompletionsMiddleware,
			Endpoint: func(c *gin.Context) {
				c.JSON(http.StatusOK, api.GenerateResponse{
					Response: "Hello!",
				})
			},
			Setup: func(t *testing.T, req *http.Request) {
				body := CompletionRequest{
					Model:  "test-model",
					Prompt: "Hello",
				}

				bodyBytes, _ := json.Marshal(body)

				req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
				req.Header.Set("Content-Type", "application/json")
			},
			Expected: func(t *testing.T, resp *httptest.ResponseRecorder) {
				assert.Equal(t, http.StatusOK, resp.Code)
				var completionResp Completion
				if err := json.NewDecoder(resp.Body).Decode(&completionResp); err != nil {
					t.Fatal(err)