|
6 | 6 | "encoding/json" |
7 | 7 | "io" |
8 | 8 | "net/http" |
| 9 | + "net/http/httptest" |
| 10 | + "net/url" |
9 | 11 | "strings" |
10 | 12 | "sync" |
11 | 13 | "testing" |
@@ -52,6 +54,86 @@ func newMockServer(mock *mockRunner) func(discover.GpuInfoList, string, *ggml.GG |
52 | 54 | } |
53 | 55 | } |
54 | 56 |
|
| 57 | +func TestGenerateChatRemote(t *testing.T) { |
| 58 | + gin.SetMode(gin.TestMode) |
| 59 | + |
| 60 | + rs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 61 | + if r.Method != http.MethodPost { |
| 62 | + t.Errorf("Expected POST request, got %s", r.Method) |
| 63 | + } |
| 64 | + if r.URL.Path != "/api/chat" { |
| 65 | + t.Errorf("Expected path '/api/chat', got %s", r.URL.Path) |
| 66 | + } |
| 67 | + |
| 68 | + w.WriteHeader(http.StatusOK) |
| 69 | + w.Header().Set("Content-Type", "application/json") |
| 70 | + resp := api.ChatResponse{ |
| 71 | + Model: "test", |
| 72 | + Done: true, |
| 73 | + DoneReason: "load", |
| 74 | + } |
| 75 | + if err := json.NewEncoder(w).Encode(&resp); err != nil { |
| 76 | + t.Fatal(err) |
| 77 | + } |
| 78 | + })) |
| 79 | + defer rs.Close() |
| 80 | + |
| 81 | + p, err := url.Parse(rs.URL) |
| 82 | + if err != nil { |
| 83 | + t.Fatal(err) |
| 84 | + } |
| 85 | + |
| 86 | + t.Setenv("OLLAMA_REMOTES", p.Hostname()) |
| 87 | + s := Server{} |
| 88 | + w := createRequest(t, s.CreateHandler, api.CreateRequest{ |
| 89 | + Model: "test-cloud", |
| 90 | + RemoteHost: rs.URL, |
| 91 | + From: "test", |
| 92 | + Info: map[string]any{ |
| 93 | + "capabilities": []string{"completion", "thinking"}, |
| 94 | + }, |
| 95 | + Stream: &stream, |
| 96 | + }) |
| 97 | + |
| 98 | + if w.Code != http.StatusOK { |
| 99 | + t.Fatalf("expected status 200, got %d", w.Code) |
| 100 | + } |
| 101 | + |
| 102 | + t.Run("missing messages", func(t *testing.T) { |
| 103 | + w := createRequest(t, s.ChatHandler, api.ChatRequest{ |
| 104 | + Model: "test-cloud", |
| 105 | + }) |
| 106 | + if w.Code != http.StatusOK { |
| 107 | + t.Errorf("expected status 200, got %d", w.Code) |
| 108 | + } |
| 109 | + |
| 110 | + var actual api.ChatResponse |
| 111 | + if err := json.NewDecoder(w.Body).Decode(&actual); err != nil { |
| 112 | + t.Fatal(err) |
| 113 | + } |
| 114 | + |
| 115 | + if actual.Model != "test-cloud" { |
| 116 | + t.Errorf("expected model test-cloud, got %s", actual.Model) |
| 117 | + } |
| 118 | + |
| 119 | + if actual.RemoteModel != "test" { |
| 120 | + t.Errorf("expected remote model test, got %s", actual.RemoteModel) |
| 121 | + } |
| 122 | + |
| 123 | + if actual.RemoteHost != rs.URL { |
| 124 | + t.Errorf("expected remote host '%s', got %s", rs.URL, actual.RemoteHost) |
| 125 | + } |
| 126 | + |
| 127 | + if !actual.Done { |
| 128 | + t.Errorf("expected done true, got false") |
| 129 | + } |
| 130 | + |
| 131 | + if actual.DoneReason != "load" { |
| 132 | + t.Errorf("expected done reason load, got %s", actual.DoneReason) |
| 133 | + } |
| 134 | + }) |
| 135 | +} |
| 136 | + |
55 | 137 | func TestGenerateChat(t *testing.T) { |
56 | 138 | gin.SetMode(gin.TestMode) |
57 | 139 |
|
|
0 commit comments