#170: Refactored the error tests

EinStack · Mar 17, 2024 · 9cad10c · 9cad10c
1 parent 29c38d4
commit 9cad10c
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 24 deletions.
diff --git a/pkg/providers/azureopenai/client_test.go b/pkg/providers/azureopenai/client_test.go
@@ -106,21 +106,19 @@ func TestDoChatRequest_ErrorResponse(t *testing.T) {
 
 	defer mockServer.Close()
 
-	// Create a new client with the mock server URL
-	client := &Client{
-		httpClient: http.DefaultClient,
-		chatURL:    mockServer.URL,
-		config:     &Config{APIKey: "dummy_key"},
-		telemetry:  telemetry.NewTelemetryMock(),
-	}
+	ctx := context.Background()
+	providerCfg := DefaultConfig()
+	clientCfg := clients.DefaultClientConfig()
+
+	providerCfg.BaseURL = mockServer.URL
+
+	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
+	require.NoError(t, err)
 
 	// Create a chat request payload
-	payload := &ChatRequest{
-		Messages: []ChatMessage{{Role: "human", Content: "Hello"}},
-	}
+	payload := schemas.NewChatFromStr("What's the dealio?")
 
-	// Call the doChatRequest function
-	_, err := client.doChatRequest(context.Background(), payload)
+	_, err = client.Chat(ctx, payload)
 
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "provider is not available")

diff --git a/pkg/providers/octoml/client_test.go b/pkg/providers/octoml/client_test.go
@@ -112,21 +112,19 @@ func TestDoChatRequest_ErrorResponse(t *testing.T) {
 	defer mockServer.Close()
 
 	// Create a new client with the mock server URL
-	client := &Client{
-		httpClient: http.DefaultClient,
-		chatURL:    mockServer.URL,
-		config:     &Config{APIKey: "dummy_key"},
-		telemetry:  telemetry.NewTelemetryMock(),
-	}
+	ctx := context.Background()
+	providerCfg := DefaultConfig()
+	clientCfg := clients.DefaultClientConfig()
+
+	providerCfg.BaseURL = mockServer.URL
+
+	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
+	require.NoError(t, err)
 
 	// Create a chat request payload
-	payload := &ChatRequest{
-		Model:    "dummy_model",
-		Messages: []ChatMessage{{Role: "human", Content: "Hello"}},
-	}
+	payload := schemas.NewChatFromStr("What's the dealio?")
 
-	// Call the doChatRequest function
-	_, err := client.doChatRequest(context.Background(), payload)
+	_, err = client.Chat(ctx, payload)
 
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "provider is not available")