#153: Fixed ollama broken tests

Author: roma-glushko

pkg/providers/ollama/client_test.go

@@ -77,20 +77,19 @@ func TestOllamaClient_ChatRequest_Non200Response(t *testing.T) {
 
 	defer mockServer.Close()
 
-	// Create a new client with the mock server URL
-	client := &Client{
-		httpClient: http.DefaultClient,
-		chatURL:    mockServer.URL,
-		config:     DefaultConfig(),
-		telemetry:  telemetry.NewTelemetryMock(),
-	}
+	providerCfg := DefaultConfig()
+	clientCfg := clients.DefaultClientConfig()
+	providerCfg.BaseURL = mockServer.URL
+
+	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
+	require.NoError(t, err)
 
 	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
 		Role:    "user",
 		Content: "What's the capital of the United Kingdom?",
 	}}}
 
-	_, err := client.Chat(context.Background(), &chatParams)
+	_, err = client.Chat(context.Background(), &chatParams)
 
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "provider is not available")
@@ -99,19 +98,28 @@ func TestOllamaClient_ChatRequest_Non200Response(t *testing.T) {
 func TestOllamaClient_ChatRequest_SuccessfulResponse(t *testing.T) {
 	// Create a mock HTTP server that returns an OK status code and a sample response
 	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		chatResponse, err := os.ReadFile(filepath.Clean("./testdata/chat.success.json"))
+		if err != nil {
+			t.Errorf("error reading cohere chat mock response: %v", err)
+		}
+
 		w.WriteHeader(http.StatusOK)
-		_, _ = w.Write([]byte(`{"response": "OK"}`))
+		w.Header().Set("Content-Type", "application/json")
+
+		_, err = w.Write(chatResponse)
+		if err != nil {
+			t.Errorf("error on sending chat response: %v", err)
+		}
 	}))
 
 	defer mockServer.Close()
 
-	// Create a new client with the mock server URL
-	client := &Client{
-		httpClient: http.DefaultClient,
-		chatURL:    mockServer.URL,
-		config:     DefaultConfig(),
-		telemetry:  telemetry.NewTelemetryMock(),
-	}
+	providerCfg := DefaultConfig()
+	clientCfg := clients.DefaultClientConfig()
+	providerCfg.BaseURL = mockServer.URL
+
+	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
+	require.NoError(t, err)
 
 	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
 		Role:    "user",
@@ -122,5 +130,6 @@ func TestOllamaClient_ChatRequest_SuccessfulResponse(t *testing.T) {
 
 	require.NoError(t, err)
 	require.NotNil(t, response)
-	require.Equal(t, "", response.ModelResponse.Message.Role)
+	require.Equal(t, "assistant", response.ModelResponse.Message.Role)
+	require.Equal(t, "London", response.ModelResponse.Message.Content)
 }

pkg/providers/ollama/testdata/chat.success.json

@@ -3,7 +3,7 @@
   "created_at": "2023-12-12T14:13:43.416799Z",
   "message": {
     "role": "assistant",
-    "content": "Hello! How are you today?"
+    "content": "London"
   },
   "done": true,
   "total_duration": 5191566416,
@@ -12,4 +12,4 @@
   "prompt_eval_duration": 383809000,
   "eval_count": 298,
   "eval_duration": 4799921000
-}
+}

pkg/providers/testing/lang.go

@@ -103,15 +103,16 @@ type ProviderMock struct {
 func NewProviderMock(modelName *string, responses []RespMock) *ProviderMock {
 	return &ProviderMock{
 		idx:              0,
-		modelName:        modelName,
 		chatResps:        &responses,
 		supportStreaming: false,
+		modelName:        modelName,
 	}
 }
 
-func NewStreamProviderMock(chatStreams []RespStreamMock) *ProviderMock {
+func NewStreamProviderMock(modelName *string, chatStreams []RespStreamMock) *ProviderMock {
 	return &ProviderMock{
 		idx:              0,
+		modelName:        modelName,
 		chatStreams:      &chatStreams,
 		supportStreaming: true,
 	}
@@ -156,7 +157,7 @@ func (c *ProviderMock) Provider() string {
 }
 
 func (c *ProviderMock) ModelName() string {
-	if c.modelName != nil {
+	if c.modelName == nil {
 		return "model_mock"
 	}
 

pkg/routers/router_test.go

@@ -262,7 +262,7 @@ func TestLangRouter_ChatStream(t *testing.T) {
 	langModels := []*providers.LanguageModel{
 		providers.NewLangModel(
 			"first",
-			ptesting.NewStreamProviderMock([]ptesting.RespStreamMock{
+			ptesting.NewStreamProviderMock(nil, []ptesting.RespStreamMock{
 				ptesting.NewRespStreamMock(&[]ptesting.RespMock{
 					{Msg: "Bill"},
 					{Msg: "Gates"},
@@ -277,7 +277,7 @@ func TestLangRouter_ChatStream(t *testing.T) {
 		),
 		providers.NewLangModel(
 			"second",
-			ptesting.NewStreamProviderMock([]ptesting.RespStreamMock{
+			ptesting.NewStreamProviderMock(nil, []ptesting.RespStreamMock{
 				ptesting.NewRespStreamMock(&[]ptesting.RespMock{
 					{Msg: "Knock"},
 					{Msg: "Knock"},
@@ -338,14 +338,14 @@ func TestLangRouter_ChatStream_FailOnFirst(t *testing.T) {
 	langModels := []*providers.LanguageModel{
 		providers.NewLangModel(
 			"first",
-			ptesting.NewStreamProviderMock(nil),
+			ptesting.NewStreamProviderMock(nil, nil),
 			budget,
 			*latConfig,
 			1,
 		),
 		providers.NewLangModel(
 			"second",
-			ptesting.NewStreamProviderMock([]ptesting.RespStreamMock{
+			ptesting.NewStreamProviderMock(nil, []ptesting.RespStreamMock{
 				ptesting.NewRespStreamMock(
 					&[]ptesting.RespMock{
 						{Msg: "Knock"},
@@ -408,7 +408,7 @@ func TestLangRouter_ChatStream_AllModelsUnavailable(t *testing.T) {
 	langModels := []*providers.LanguageModel{
 		providers.NewLangModel(
 			"first",
-			ptesting.NewStreamProviderMock([]ptesting.RespStreamMock{
+			ptesting.NewStreamProviderMock(nil, []ptesting.RespStreamMock{
 				ptesting.NewRespStreamMock(&[]ptesting.RespMock{
 					{Err: clients.ErrProviderUnavailable},
 				}),
@@ -419,7 +419,7 @@ func TestLangRouter_ChatStream_AllModelsUnavailable(t *testing.T) {
 		),
 		providers.NewLangModel(
 			"second",
-			ptesting.NewStreamProviderMock([]ptesting.RespStreamMock{
+			ptesting.NewStreamProviderMock(nil, []ptesting.RespStreamMock{
 				ptesting.NewRespStreamMock(&[]ptesting.RespMock{
 					{Err: clients.ErrProviderUnavailable},
 				}),