Skip to content

Commit

Permalink
added unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
AyushSawant18588 committed Nov 23, 2024
1 parent 5b15527 commit 88775cf
Showing 1 changed file with 33 additions and 0 deletions.
33 changes: 33 additions & 0 deletions embeddings_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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(),
Expand Down

0 comments on commit 88775cf

Please sign in to comment.