You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/tests/models/ada.py

77 lines
2.3 KiB

1 year ago
# test_embeddings.py
import pytest
import openai
from unittest.mock import patch
from swarms.models.simple_ada import (
get_ada_embeddings,
) # Adjust this import path to your project structure
1 year ago
from os import getenv
from dotenv import load_dotenv
load_dotenv()
1 year ago
# Fixture for test texts
@pytest.fixture
def test_texts():
return [
"Hello World",
"This is a test string with newline\ncharacters",
"A quick brown fox jumps over the lazy dog",
]
1 year ago
# Basic Test
def test_get_ada_embeddings_basic(test_texts):
with patch("openai.Embedding.create") as mock_create:
1 year ago
# Mocking the OpenAI API call
mock_create.return_value = {"data": [{"embedding": [0.1, 0.2, 0.3]}]}
1 year ago
for text in test_texts:
embedding = get_ada_embeddings(text)
assert embedding == [
0.1,
0.2,
0.3,
], "Embedding does not match expected output"
mock_create.assert_called_with(
input=[text.replace("\n", " ")], model="text-embedding-ada-002"
)
1 year ago
# Parameterized Test
@pytest.mark.parametrize(
"text, model, expected_call_model",
[
("Hello World", "text-embedding-ada-002", "text-embedding-ada-002"),
("Hello World", "text-embedding-ada-001", "text-embedding-ada-001"),
],
)
def test_get_ada_embeddings_models(text, model, expected_call_model):
with patch("openai.Embedding.create") as mock_create:
mock_create.return_value = {"data": [{"embedding": [0.1, 0.2, 0.3]}]}
1 year ago
_ = get_ada_embeddings(text, model=model)
mock_create.assert_called_with(input=[text], model=expected_call_model)
1 year ago
# Exception Test
def test_get_ada_embeddings_exception():
with patch("openai.Embedding.create") as mock_create:
1 year ago
mock_create.side_effect = openai.error.OpenAIError("Test error")
with pytest.raises(openai.error.OpenAIError):
get_ada_embeddings("Some text")
1 year ago
# Tests for environment variable loading
def test_env_var_loading(monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "testkey123")
with patch("openai.Embedding.create"):
assert (
getenv("OPENAI_API_KEY") == "testkey123"
), "Environment variable for API key is not set correctly"
1 year ago
# ... more tests to cover other aspects such as different input types, large inputs, invalid inputs, etc.