#153: Updating tests to feed chat params instead of request schema

EinStack · Jun 4, 2024 · 85e7646 · 85e7646
1 parent 942a9d9
commit 85e7646
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 22 deletions.
diff --git a/pkg/providers/azureopenai/chat_stream_test.go b/pkg/providers/azureopenai/chat_stream_test.go
@@ -71,8 +71,12 @@ func TestAzureOpenAIClient_ChatStreamRequest(t *testing.T) {
 			client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
 			require.NoError(t, err)
 
-			req := schemas.NewChatStreamFromStr("What's the capital of the United Kingdom?")
-			stream, err := client.ChatStream(ctx, req)
+			chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
+				Role:    "user",
+				Content: "What's the capital of the United Kingdom?",
+			}}}
+
+			stream, err := client.ChatStream(ctx, &chatParams)
 			require.NoError(t, err)
 
 			err = stream.Open()

diff --git a/pkg/providers/bedrock/client_test.go b/pkg/providers/bedrock/client_test.go
@@ -61,12 +61,12 @@ func TestBedrockClient_ChatRequest(t *testing.T) {
 	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
 	require.NoError(t, err)
 
-	request := schemas.ChatRequest{Message: schemas.ChatMessage{
+	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
 		Role:    "user",
 		Content: "What's the biggest animal?",
-	}}
+	}}}
 
-	response, err := client.Chat(ctx, &request)
+	response, err := client.Chat(ctx, &chatParams)
 
 	responseString := fmt.Sprintf("%+v", response)
 	// errString := fmt.Sprintf("%+v", err)

diff --git a/pkg/providers/octoml/client_test.go b/pkg/providers/octoml/client_test.go
@@ -55,12 +55,12 @@ func TestOctoMLClient_ChatRequest(t *testing.T) {
 	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
 	require.NoError(t, err)
 
-	request := schemas.ChatRequest{Message: schemas.ChatMessage{
+	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
 		Role:    "human",
 		Content: "What's the biggest animal?",
-	}}
+	}}}
 
-	response, err := client.Chat(ctx, &request)
+	response, err := client.Chat(ctx, &chatParams)
 	require.NoError(t, err)
 
 	require.Equal(t, providerCfg.ModelName, response.ModelName)
@@ -88,15 +88,13 @@ func TestOctoMLClient_Chat_Error(t *testing.T) {
 	require.NoError(t, err)
 
 	// Create a chat request
-	request := schemas.ChatRequest{
-		Message: schemas.ChatMessage{
-			Role:    "human",
-			Content: "What's the biggest animal?",
-		},
-	}
+	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
+		Role:    "human",
+		Content: "What's the biggest animal?",
+	}}}
 
 	// Call the Chat function
-	_, err = client.Chat(ctx, &request)
+	_, err = client.Chat(ctx, &chatParams)
 
 	// Check the error
 	require.Error(t, err)
@@ -122,9 +120,12 @@ func TestDoChatRequest_ErrorResponse(t *testing.T) {
 	require.NoError(t, err)
 
 	// Create a chat request payload
-	payload := schemas.NewChatFromStr("What's the dealio?")
+	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
+		Role:    "user",
+		Content: "What's the dealeo?",
+	}}}
 
-	_, err = client.Chat(ctx, payload)
+	_, err = client.Chat(ctx, &chatParams)
 
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "provider is not available")

diff --git a/pkg/providers/ollama/client_test.go b/pkg/providers/ollama/client_test.go
@@ -56,12 +56,12 @@ func TestOllamaClient_ChatRequest(t *testing.T) {
 	client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
 	require.NoError(t, err)
 
-	request := schemas.ChatRequest{Message: schemas.ChatMessage{
+	chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
 		Role:    "user",
 		Content: "What's the biggest animal?",
-	}}
+	}}}
 
-	_, err = client.Chat(ctx, &request)
+	_, err = client.Chat(ctx, &chatParams)
 
 	// require.NoError(t, err)
 

diff --git a/pkg/providers/openai/chat_stream_test.go b/pkg/providers/openai/chat_stream_test.go
@@ -71,8 +71,12 @@ func TestOpenAIClient_ChatStreamRequest(t *testing.T) {
 			client, err := NewClient(providerCfg, clientCfg, telemetry.NewTelemetryMock())
 			require.NoError(t, err)
 
-			req := schemas.NewChatStreamFromStr("What's the capital of the United Kingdom?")
-			stream, err := client.ChatStream(ctx, req)
+			chatParams := schemas.ChatParams{Messages: []schemas.ChatMessage{{
+				Role:    "user",
+				Content: "What's the capital of the United Kingdom?",
+			}}}
+
+			stream, err := client.ChatStream(ctx, &chatParams)
 			require.NoError(t, err)
 
 			err = stream.Open()