From fd48443b6c7b672ca2506fa75f1c68abca6f11d6 Mon Sep 17 00:00:00 2001 From: Daemon <109057945+Daethyra@users.noreply.github.com> Date: Wed, 4 Oct 2023 14:18:08 -0700 Subject: [PATCH] Updated test.py --- Auto-Embedder/test.py | 162 ++++++++++++++++++++++++++---------------- 1 file changed, 100 insertions(+), 62 deletions(-) diff --git a/Auto-Embedder/test.py b/Auto-Embedder/test.py index 517312a..aadc19d 100644 --- a/Auto-Embedder/test.py +++ b/Auto-Embedder/test.py @@ -1,71 +1,109 @@ -import unittest -from unittest.mock import patch -import os -import logging + +import pytest import asyncio -from dotenv import load_dotenv -from datetime import datetime +import os +from unittest.mock import patch, Mock from typing import Dict, Union, List -import openai -import pinecone -import backoff +from pinembed import EnvConfig, OpenAIHandler, PineconeHandler, DataStreamHandler -# Load environment variables from .env file -load_dotenv() - -# Initialize logging -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - -class EnvConfig: - """Class for handling environment variables and API keys.""" - - def __init__(self) -> None: - """Initialize environment variables.""" - self.openai_key: str = os.getenv("OPENAI_API_KEY") - self.pinecone_key: str = os.getenv("PINECONE_API_KEY") - self.pinecone_environment: str = os.getenv("PINECONE_ENVIRONMENT") - self.pinecone_environment: str = os.getenv("PINEDEX") +@pytest.fixture +def mock_env_config(): + """Fixture for setting up a mock environment configuration. + + Mocks os.getenv to return a test value and initializes EnvConfig. + + Returns: + EnvConfig: Mocked environment configuration + """ + with patch('os.getenv', return_value="test_value"): + config = EnvConfig() + return config -class OpenAIHandler: - """Class for handling OpenAI operations.""" +@pytest.mark.parametrize("env_value, expected_value", [("test_value", "test_value"), (None, None)]) +def test_EnvConfig_init(env_value, expected_value, mock_env_config): + """Test initialization of EnvConfig. + + Tests if the EnvConfig is correctly initialized with environment variables. + + Args: + env_value (str or None): Mock environment variable value + expected_value (str or None): Expected value for EnvConfig attributes + mock_env_config (EnvConfig): Mocked environment configuration + """ + assert mock_env_config.openai_key == expected_value + assert mock_env_config.pinecone_key == expected_value - def __init__(self, config: EnvConfig) -> None: - """Initialize OpenAI API key.""" - openai.api_key = config.openai_key - - @backoff.on_exception(backoff.expo, Exception, max_tries=3) - async def create_embedding(self, input_text: str) -> Dict[str, Union[int, List[float]]]: - """ - Create an embedding using OpenAI. +@pytest.mark.asyncio +@pytest.mark.parallel +async def test_OpenAIHandler_create_embedding(mock_env_config): + """Asynchronous test for creating embeddings via OpenAIHandler. + + Tests if OpenAIHandler.create_embedding method correctly returns mock response. + + Args: + mock_env_config (EnvConfig): Mocked environment configuration + """ + handler = OpenAIHandler(mock_env_config) + mock_response = {"id": 1, "values": [0.1, 0.2, 0.3]} + + with patch.object(handler.openai.Embedding, 'create', return_value=mock_response): + response = await handler.create_embedding("test_text") - Parameters: - input_text (str): The text to be embedded. - - Returns: - Dict[str, Union[int, List[float]]]: The embedding response. - """ - response = openai.Embedding.create( - model="text-embedding-ada-002",engine="ada", - text=input_text, - ) - return response + assert response == mock_response -# Create test class -class TestOpenAIHandler(unittest.TestCase): - # Set up test environment - def setUp(self): - self.config = EnvConfig() - self.openai_handler = OpenAIHandler(self.config) +@pytest.mark.parallel +def test_PineconeHandler_init(mock_env_config): + """Test initialization of PineconeHandler. + + Tests if PineconeHandler is correctly initialized with environment variables. + + Args: + mock_env_config (EnvConfig): Mocked environment configuration + """ + handler = PineconeHandler(mock_env_config) + handler.pinecone.init.assert_called_with(api_key="test_value", environment="test_value") + assert handler.index_name == "test_value" - # Test create_embedding method - @patch('openai.Embedding.create') - def test_create_embedding(self, mock_create): - input_text = 'This is a test' - expected_response = {'id': 12345, 'embedding': [1.0, 2.0, 3.0]} - mock_create.return_value = expected_response - response = self.openai_handler.create_embedding(input_text) - self.assertEqual(response, expected_response) +@pytest.mark.asyncio +@pytest.mark.parallel +async def test_PineconeHandler_upload_embedding(mock_env_config): + """Asynchronous test for uploading embeddings via PineconeHandler. + + Tests if PineconeHandler.upload_embedding method correctly calls pinecone.Index.upsert. + + Args: + mock_env_config (EnvConfig): Mocked environment configuration + """ + handler = PineconeHandler(mock_env_config) + mock_embedding = { + "id": "1", + "values": [0.1, 0.2, 0.3], + "metadata": {}, + "sparse_values": {} + } + + with patch.object(handler.pinecone.Index, 'upsert', return_value=None): + await handler.upload_embedding(mock_embedding) + + handler.pinecone.Index.assert_called_with("test_value") -if __name__ == "__main__": - unittest.main() \ No newline at end of file +@pytest.mark.asyncio +@pytest.mark.parallel +async def test_DataStreamHandler_process_data(mock_env_config): + """Asynchronous test for processing data via DataStreamHandler. + + Tests if DataStreamHandler.process_data method correctly calls methods of OpenAIHandler and PineconeHandler. + + Args: + mock_env_config (EnvConfig): Mocked environment configuration + """ + openai_handler = OpenAIHandler(mock_env_config) + pinecone_handler = PineconeHandler(mock_env_config) + handler = DataStreamHandler(openai_handler, pinecone_handler) + + mock_data = "test_data" + mock_embedding = {"id": 1, "values": [0.1, 0.2, 0.3]} + + with patch.object(OpenAIHandler, 'create_embedding', return_value=mock_embedding): + with patch.object(PineconeHandler, 'upload_embedding', return_value=None): + await handler.process_data(mock_data)