Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
client_test.go 8.79 KB
Newer Older
Michael Yang's avatar
client: fix trailing slash  
Michael Yang committed
1
2 
package api
Patrick Devine's avatar
better checking for OLLAMA_HOST variable (#3661)  
Patrick Devine committed
3 
import (
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
4
5
6
7
8
9 
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
Patrick Devine's avatar
better checking for OLLAMA_HOST variable (#3661)  
Patrick Devine committed
10
11 
	"testing"
)
Michael Yang's avatar
client: fix trailing slash  
Michael Yang committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51 

func TestClientFromEnvironment(t *testing.T) {
	type testCase struct {
		value  string
		expect string
		err    error
	}

	testCases := map[string]*testCase{
		"empty":                      {value: "", expect: "http://127.0.0.1:11434"},
		"only address":               {value: "1.2.3.4", expect: "http://1.2.3.4:11434"},
		"only port":                  {value: ":1234", expect: "http://:1234"},
		"address and port":           {value: "1.2.3.4:1234", expect: "http://1.2.3.4:1234"},
		"scheme http and address":    {value: "http://1.2.3.4", expect: "http://1.2.3.4:80"},
		"scheme https and address":   {value: "https://1.2.3.4", expect: "https://1.2.3.4:443"},
		"scheme, address, and port":  {value: "https://1.2.3.4:1234", expect: "https://1.2.3.4:1234"},
		"hostname":                   {value: "example.com", expect: "http://example.com:11434"},
		"hostname and port":          {value: "example.com:1234", expect: "http://example.com:1234"},
		"scheme http and hostname":   {value: "http://example.com", expect: "http://example.com:80"},
		"scheme https and hostname":  {value: "https://example.com", expect: "https://example.com:443"},
		"scheme, hostname, and port": {value: "https://example.com:1234", expect: "https://example.com:1234"},
		"trailing slash":             {value: "example.com/", expect: "http://example.com:11434"},
		"trailing slash port":        {value: "example.com:1234/", expect: "http://example.com:1234"},
	}

	for k, v := range testCases {
		t.Run(k, func(t *testing.T) {
			t.Setenv("OLLAMA_HOST", v.value)

			client, err := ClientFromEnvironment()
			if err != v.err {
				t.Fatalf("expected %s, got %s", v.err, err)
			}

			if client.base.String() != v.expect {
				t.Fatalf("expected %s, got %s", v.expect, client.base.String())
			}
		})
	}
}
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
52
53
54
55
56
57 

// testError represents an internal error type with status code and message
// this is used since the error response from the server is not a standard error struct
type testError struct {
	message    string
	statusCode int
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
58 
	raw        bool // if true, write message as-is instead of JSON encoding
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92 
}

func (e testError) Error() string {
	return e.message
}

func TestClientStream(t *testing.T) {
	testCases := []struct {
		name      string
		responses []any
		wantErr   string
	}{
		{
			name: "immediate error response",
			responses: []any{
				testError{
					message:    "test error message",
					statusCode: http.StatusBadRequest,
				},
			},
			wantErr: "test error message",
		},
		{
			name: "error after successful chunks, ok response",
			responses: []any{
				ChatResponse{Message: Message{Content: "partial response 1"}},
				ChatResponse{Message: Message{Content: "partial response 2"}},
				testError{
					message:    "mid-stream error",
					statusCode: http.StatusOK,
				},
			},
			wantErr: "mid-stream error",
		},
Bruce MacDonald's avatar
api: fix unreachable status err (#11423)  
Bruce MacDonald committed
93
94
95
96
97
98
99
100
101
102 
		{
			name: "http status error takes precedence over general error",
			responses: []any{
				testError{
					message:    "custom error message",
					statusCode: http.StatusInternalServerError,
				},
			},
			wantErr: "500",
		},
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
103
104
105
106
107
108
109
110
111
112
113
114 
		{
			name: "successful stream completion",
			responses: []any{
				ChatResponse{Message: Message{Content: "chunk 1"}},
				ChatResponse{Message: Message{Content: "chunk 2"}},
				ChatResponse{
					Message:    Message{Content: "final chunk"},
					Done:       true,
					DoneReason: "stop",
				},
			},
		},
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128 
		{
			name: "plain text error response",
			responses: []any{
				"internal server error",
			},
			wantErr: "internal server error",
		},
		{
			name: "HTML error page",
			responses: []any{
				"<html><body>404 Not Found</body></html>",
			},
			wantErr: "404 Not Found",
		},
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152 
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				flusher, ok := w.(http.Flusher)
				if !ok {
					t.Fatal("expected http.Flusher")
				}

				w.Header().Set("Content-Type", "application/x-ndjson")

				for _, resp := range tc.responses {
					if errResp, ok := resp.(testError); ok {
						w.WriteHeader(errResp.statusCode)
						err := json.NewEncoder(w).Encode(map[string]string{
							"error": errResp.message,
						})
						if err != nil {
							t.Fatal("failed to encode error response:", err)
						}
						return
					}
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
153
154
155
156
157
158 
					if str, ok := resp.(string); ok {
						fmt.Fprintln(w, str)
						flusher.Flush()
						continue
					}
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
159
160
161
162
163
164
165
166
167
168
169 
					if err := json.NewEncoder(w).Encode(resp); err != nil {
						t.Fatalf("failed to encode response: %v", err)
					}
					flusher.Flush()
				}
			}))
			defer ts.Close()

			client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)

			var receivedChunks []ChatResponse
Michael Yang's avatar
lint: enable usetesting, disable tenv (#10594)  
Michael Yang committed
170 
			err := client.stream(t.Context(), http.MethodPost, "/v1/chat", nil, func(chunk []byte) error {
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196 
				var resp ChatResponse
				if err := json.Unmarshal(chunk, &resp); err != nil {
					return fmt.Errorf("failed to unmarshal chunk: %w", err)
				}
				receivedChunks = append(receivedChunks, resp)
				return nil
			})

			if tc.wantErr != "" {
				if err == nil {
					t.Fatal("expected error but got nil")
				}
				if !strings.Contains(err.Error(), tc.wantErr) {
					t.Errorf("expected error containing %q, got %v", tc.wantErr, err)
				}
				return
			}
			if err != nil {
				t.Errorf("unexpected error: %v", err)
			}
		})
	}
}

func TestClientDo(t *testing.T) {
	testCases := []struct {
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
197
198
199
200 
		name           string
		response       any
		wantErr        string
		wantStatusCode int
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
201
202
203
204
205
206
207 
	}{
		{
			name: "immediate error response",
			response: testError{
				message:    "test error message",
				statusCode: http.StatusBadRequest,
			},
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
208
209 
			wantErr:        "test error message",
			wantStatusCode: http.StatusBadRequest,
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
210
211
212
213
214
215
216 
		},
		{
			name: "server error response",
			response: testError{
				message:    "internal error",
				statusCode: http.StatusInternalServerError,
			},
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
217
218 
			wantErr:        "internal error",
			wantStatusCode: http.StatusInternalServerError,
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
219
220
221
222
223
224
225
226
227
228
229 
		},
		{
			name: "successful response",
			response: struct {
				ID      string `json:"id"`
				Success bool   `json:"success"`
			}{
				ID:      "msg_123",
				Success: true,
			},
		},
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249 
		{
			name: "plain text error response",
			response: testError{
				message:    "internal server error",
				statusCode: http.StatusInternalServerError,
				raw:        true,
			},
			wantErr:        "internal server error",
			wantStatusCode: http.StatusInternalServerError,
		},
		{
			name: "HTML error page",
			response: testError{
				message:    "<html><body>404 Not Found</body></html>",
				statusCode: http.StatusNotFound,
				raw:        true,
			},
			wantErr:        "<html><body>404 Not Found</body></html>",
			wantStatusCode: http.StatusNotFound,
		},
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
250
251
252
253
254
255
256 
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if errResp, ok := tc.response.(testError); ok {
					w.WriteHeader(errResp.statusCode)
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
257
258
259
260
261
262
263
264
265
266 
					if !errResp.raw {
						err := json.NewEncoder(w).Encode(map[string]string{
							"error": errResp.message,
						})
						if err != nil {
							t.Fatal("failed to encode error response:", err)
						}
					} else {
						// Write raw message (simulates non-JSON error responses)
						fmt.Fprint(w, errResp.message)
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283 
					}
					return
				}

				w.Header().Set("Content-Type", "application/json")
				if err := json.NewEncoder(w).Encode(tc.response); err != nil {
					t.Fatalf("failed to encode response: %v", err)
				}
			}))
			defer ts.Close()

			client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)

			var resp struct {
				ID      string `json:"id"`
				Success bool   `json:"success"`
			}
Michael Yang's avatar
lint: enable usetesting, disable tenv (#10594)  
Michael Yang committed
284 
			err := client.do(t.Context(), http.MethodPost, "/v1/messages", nil, &resp)
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
285
286
287
288
289
290
291
292 

			if tc.wantErr != "" {
				if err == nil {
					t.Fatalf("got nil, want error %q", tc.wantErr)
				}
				if err.Error() != tc.wantErr {
					t.Errorf("error message mismatch: got %q, want %q", err.Error(), tc.wantErr)
				}
Bruce MacDonald's avatar
api/client: handle non-json streaming errors (#13007)  
Bruce MacDonald committed
293
294
295
296
297
298
299
300
301 
				if tc.wantStatusCode != 0 {
					if statusErr, ok := err.(StatusError); ok {
						if statusErr.StatusCode != tc.wantStatusCode {
							t.Errorf("status code mismatch: got %d, want %d", statusErr.StatusCode, tc.wantStatusCode)
						}
					} else {
						t.Errorf("expected StatusError, got %T", err)
					}
				}
Bruce MacDonald's avatar
api: document client stream behavior with a test (#8996)  
Bruce MacDonald committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322 
				return
			}

			if err != nil {
				t.Fatalf("got error %q, want nil", err)
			}

			if expectedResp, ok := tc.response.(struct {
				ID      string `json:"id"`
				Success bool   `json:"success"`
			}); ok {
				if resp.ID != expectedResp.ID {
					t.Errorf("response ID mismatch: got %q, want %q", resp.ID, expectedResp.ID)
				}
				if resp.Success != expectedResp.Success {
					t.Errorf("response Success mismatch: got %v, want %v", resp.Success, expectedResp.Success)
				}
			}
		})
	}
}