From 88775cf9bff49323ef76ac7673dcdf7e0c393302 Mon Sep 17 00:00:00 2001
From: ayush <ayush.sawant@nutanix.com>
Date: Sat, 23 Nov 2024 16:43:22 +0000
Subject: [PATCH] added unit tests

---
 embeddings_test.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/embeddings_test.go b/embeddings_test.go
index 438978169..3eca4c31c 100644
--- a/embeddings_test.go
+++ b/embeddings_test.go
@@ -51,6 +51,24 @@ func TestEmbedding(t *testing.T) {
 			t.Fatalf("Expected embedding request to contain model field")
 		}
 
+		// test embedding request with strings and extra_body param
+		embeddingReqWithExtraBody := openai.EmbeddingRequest{
+			Input: []string{
+				"The food was delicious and the waiter",
+				"Other examples of embedding request",
+			},
+			Model: model,
+			ExtraBody: map[string]any{
+				"input_type": "query",
+				"truncate":   "NONE",
+			},
+		}
+		marshaled, err = json.Marshal(embeddingReqWithExtraBody)
+		checks.NoError(t, err, "Could not marshal embedding request")
+		if !bytes.Contains(marshaled, []byte(`"model":"`+model+`"`)) {
+			t.Fatalf("Expected embedding request to contain model field")
+		}
+
 		// test embedding request with strings
 		embeddingReqStrings := openai.EmbeddingRequestStrings{
 			Input: []string{
@@ -124,6 +142,21 @@ func TestEmbeddingEndpoint(t *testing.T) {
 		t.Errorf("Expected %#v embeddings, got %#v", sampleEmbeddings, res.Data)
 	}
 
+	// test create embeddings with strings (ExtraBody in request)
+	res, err = client.CreateEmbeddings(
+		context.Background(),
+		openai.EmbeddingRequest{
+			ExtraBody: map[string]any{
+				"input_type": "query",
+				"truncate":   "NONE",
+			},
+		},
+	)
+	checks.NoError(t, err, "CreateEmbeddings error")
+	if !reflect.DeepEqual(res.Data, sampleEmbeddings) {
+		t.Errorf("Expected %#v embeddings, got %#v", sampleEmbeddings, res.Data)
+	}
+
 	// test create embeddings with strings (simple embedding request)
 	res, err = client.CreateEmbeddings(
 		context.Background(),