From 0e49d09744c3f3340288981741ccbfe482472f4b Mon Sep 17 00:00:00 2001 From: Kye Date: Fri, 22 Dec 2023 21:44:40 -0500 Subject: [PATCH] [CLEANUP] [TESTS] --- tests/models/test_LLM.py | 56 ----- tests/models/test_dalle3.py | 454 ------------------------------------ 2 files changed, 510 deletions(-) delete mode 100644 tests/models/test_LLM.py delete mode 100644 tests/models/test_dalle3.py diff --git a/tests/models/test_LLM.py b/tests/models/test_LLM.py deleted file mode 100644 index 44bc29c9..00000000 --- a/tests/models/test_LLM.py +++ /dev/null @@ -1,56 +0,0 @@ -import unittest -import os -from unittest.mock import patch -from langchain import HuggingFaceHub -from langchain.chat_models import ChatOpenAI - -from swarms.models.llm import LLM - - -class TestLLM(unittest.TestCase): - @patch.object(HuggingFaceHub, "__init__", return_value=None) - @patch.object(ChatOpenAI, "__init__", return_value=None) - def setUp(self, mock_hf_init, mock_openai_init): - self.llm_openai = LLM(openai_api_key="mock_openai_key") - self.llm_hf = LLM( - hf_repo_id="mock_repo_id", hf_api_token="mock_hf_token" - ) - self.prompt = "Who won the FIFA World Cup in 1998?" - - def test_init(self): - self.assertEqual( - self.llm_openai.openai_api_key, "mock_openai_key" - ) - self.assertEqual(self.llm_hf.hf_repo_id, "mock_repo_id") - self.assertEqual(self.llm_hf.hf_api_token, "mock_hf_token") - - @patch.object(HuggingFaceHub, "run", return_value="France") - @patch.object(ChatOpenAI, "run", return_value="France") - def test_run(self, mock_hf_run, mock_openai_run): - result_openai = self.llm_openai.run(self.prompt) - mock_openai_run.assert_called_once() - self.assertEqual(result_openai, "France") - - result_hf = self.llm_hf.run(self.prompt) - mock_hf_run.assert_called_once() - self.assertEqual(result_hf, "France") - - def test_error_on_no_keys(self): - with self.assertRaises(ValueError): - LLM() - - @patch.object(os, "environ", {}) - def test_error_on_missing_hf_token(self): - with self.assertRaises(ValueError): - LLM(hf_repo_id="mock_repo_id") - - @patch.dict( - os.environ, {"HUGGINGFACEHUB_API_TOKEN": "mock_hf_token"} - ) - def test_hf_token_from_env(self): - llm = LLM(hf_repo_id="mock_repo_id") - self.assertEqual(llm.hf_api_token, "mock_hf_token") - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/models/test_dalle3.py b/tests/models/test_dalle3.py deleted file mode 100644 index 00ba7bc9..00000000 --- a/tests/models/test_dalle3.py +++ /dev/null @@ -1,454 +0,0 @@ -import os -from unittest.mock import Mock - -import pytest -from openai import OpenAIError -from PIL import Image -from termcolor import colored - -from swarms.models.dalle3 import Dalle3 - - -# Mocking the OpenAI client to avoid making actual API calls during testing -@pytest.fixture -def mock_openai_client(): - return Mock() - - -@pytest.fixture -def dalle3(mock_openai_client): - return Dalle3(client=mock_openai_client) - - -def test_dalle3_call_success(dalle3, mock_openai_client): - # Arrange - task = "A painting of a dog" - expected_img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - mock_openai_client.images.generate.return_value = Mock( - data=[Mock(url=expected_img_url)] - ) - - # Act - img_url = dalle3(task) - - # Assert - assert img_url == expected_img_url - mock_openai_client.images.generate.assert_called_once_with( - prompt=task, n=4 - ) - - -def test_dalle3_call_failure(dalle3, mock_openai_client, capsys): - # Arrange - task = "Invalid task" - expected_error_message = "Error running Dalle3: API Error" - - # Mocking OpenAIError - mock_openai_client.images.generate.side_effect = OpenAIError( - expected_error_message, - http_status=500, - error="Internal Server Error", - ) - - # Act and assert - with pytest.raises(OpenAIError) as excinfo: - dalle3(task) - - assert str(excinfo.value) == expected_error_message - mock_openai_client.images.generate.assert_called_once_with( - prompt=task, n=4 - ) - - # Ensure the error message is printed in red - captured = capsys.readouterr() - assert colored(expected_error_message, "red") in captured.out - - -def test_dalle3_create_variations_success(dalle3, mock_openai_client): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - expected_variation_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_02ABCDE.png" - mock_openai_client.images.create_variation.return_value = Mock( - data=[Mock(url=expected_variation_url)] - ) - - # Act - variation_img_url = dalle3.create_variations(img_url) - - # Assert - assert variation_img_url == expected_variation_url - mock_openai_client.images.create_variation.assert_called_once() - _, kwargs = mock_openai_client.images.create_variation.call_args - assert kwargs["img"] is not None - assert kwargs["n"] == 4 - assert kwargs["size"] == "1024x1024" - - -def test_dalle3_create_variations_failure( - dalle3, mock_openai_client, capsys -): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - expected_error_message = "Error running Dalle3: API Error" - - # Mocking OpenAIError - mock_openai_client.images.create_variation.side_effect = ( - OpenAIError( - expected_error_message, - http_status=500, - error="Internal Server Error", - ) - ) - - # Act and assert - with pytest.raises(OpenAIError) as excinfo: - dalle3.create_variations(img_url) - - assert str(excinfo.value) == expected_error_message - mock_openai_client.images.create_variation.assert_called_once() - - # Ensure the error message is printed in red - captured = capsys.readouterr() - assert colored(expected_error_message, "red") in captured.out - - -def test_dalle3_read_img(): - # Arrange - img_path = "test_image.png" - img = Image.new("RGB", (512, 512)) - - # Save the image temporarily - img.save(img_path) - - # Act - dalle3 = Dalle3() - img_loaded = dalle3.read_img(img_path) - - # Assert - assert isinstance(img_loaded, Image.Image) - - # Clean up - os.remove(img_path) - - -def test_dalle3_set_width_height(): - # Arrange - img = Image.new("RGB", (512, 512)) - width = 256 - height = 256 - - # Act - dalle3 = Dalle3() - img_resized = dalle3.set_width_height(img, width, height) - - # Assert - assert img_resized.size == (width, height) - - -def test_dalle3_convert_to_bytesio(): - # Arrange - img = Image.new("RGB", (512, 512)) - expected_format = "PNG" - - # Act - dalle3 = Dalle3() - img_bytes = dalle3.convert_to_bytesio(img, format=expected_format) - - # Assert - assert isinstance(img_bytes, bytes) - assert img_bytes.startswith(b"\x89PNG") - - -def test_dalle3_call_multiple_times(dalle3, mock_openai_client): - # Arrange - task = "A painting of a dog" - expected_img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - mock_openai_client.images.generate.return_value = Mock( - data=[Mock(url=expected_img_url)] - ) - - # Act - img_url1 = dalle3(task) - img_url2 = dalle3(task) - - # Assert - assert img_url1 == expected_img_url - assert img_url2 == expected_img_url - assert mock_openai_client.images.generate.call_count == 2 - - -def test_dalle3_call_with_large_input(dalle3, mock_openai_client): - # Arrange - task = "A" * 2048 # Input longer than API's limit - expected_error_message = "Error running Dalle3: API Error" - mock_openai_client.images.generate.side_effect = OpenAIError( - expected_error_message, - http_status=500, - error="Internal Server Error", - ) - - # Act and assert - with pytest.raises(OpenAIError) as excinfo: - dalle3(task) - - assert str(excinfo.value) == expected_error_message - - -def test_dalle3_create_variations_with_invalid_image_url( - dalle3, mock_openai_client -): - # Arrange - img_url = "https://invalid-image-url.com" - expected_error_message = "Error running Dalle3: Invalid image URL" - - # Act and assert - with pytest.raises(ValueError) as excinfo: - dalle3.create_variations(img_url) - - assert str(excinfo.value) == expected_error_message - - -def test_dalle3_set_width_height_invalid_dimensions(dalle3): - # Arrange - img = dalle3.read_img("test_image.png") - width = 0 - height = -1 - - # Act and assert - with pytest.raises(ValueError): - dalle3.set_width_height(img, width, height) - - -def test_dalle3_convert_to_bytesio_invalid_format(dalle3): - # Arrange - img = dalle3.read_img("test_image.png") - invalid_format = "invalid_format" - - # Act and assert - with pytest.raises(ValueError): - dalle3.convert_to_bytesio(img, format=invalid_format) - - -def test_dalle3_call_with_retry(dalle3, mock_openai_client): - # Arrange - task = "A painting of a dog" - expected_img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - - # Simulate a retry scenario - mock_openai_client.images.generate.side_effect = [ - OpenAIError( - "Temporary error", - http_status=500, - error="Internal Server Error", - ), - Mock(data=[Mock(url=expected_img_url)]), - ] - - # Act - img_url = dalle3(task) - - # Assert - assert img_url == expected_img_url - assert mock_openai_client.images.generate.call_count == 2 - - -def test_dalle3_create_variations_with_retry( - dalle3, mock_openai_client -): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - expected_variation_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_02ABCDE.png" - - # Simulate a retry scenario - mock_openai_client.images.create_variation.side_effect = [ - OpenAIError( - "Temporary error", - http_status=500, - error="Internal Server Error", - ), - Mock(data=[Mock(url=expected_variation_url)]), - ] - - # Act - variation_img_url = dalle3.create_variations(img_url) - - # Assert - assert variation_img_url == expected_variation_url - assert mock_openai_client.images.create_variation.call_count == 2 - - -def test_dalle3_call_exception_logging( - dalle3, mock_openai_client, capsys -): - # Arrange - task = "A painting of a dog" - expected_error_message = "Error running Dalle3: API Error" - - # Mocking OpenAIError - mock_openai_client.images.generate.side_effect = OpenAIError( - expected_error_message, - http_status=500, - error="Internal Server Error", - ) - - # Act - with pytest.raises(OpenAIError): - dalle3(task) - - # Assert that the error message is logged - captured = capsys.readouterr() - assert expected_error_message in captured.err - - -def test_dalle3_create_variations_exception_logging( - dalle3, mock_openai_client, capsys -): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - expected_error_message = "Error running Dalle3: API Error" - - # Mocking OpenAIError - mock_openai_client.images.create_variation.side_effect = ( - OpenAIError( - expected_error_message, - http_status=500, - error="Internal Server Error", - ) - ) - - # Act - with pytest.raises(OpenAIError): - dalle3.create_variations(img_url) - - # Assert that the error message is logged - captured = capsys.readouterr() - assert expected_error_message in captured.err - - -def test_dalle3_read_img_invalid_path(dalle3): - # Arrange - invalid_img_path = "invalid_image_path.png" - - # Act and assert - with pytest.raises(FileNotFoundError): - dalle3.read_img(invalid_img_path) - - -def test_dalle3_call_no_api_key(): - # Arrange - task = "A painting of a dog" - dalle3 = Dalle3(api_key=None) - expected_error_message = ( - "Error running Dalle3: API Key is missing" - ) - - # Act and assert - with pytest.raises(ValueError) as excinfo: - dalle3(task) - - assert str(excinfo.value) == expected_error_message - - -def test_dalle3_create_variations_no_api_key(): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - dalle3 = Dalle3(api_key=None) - expected_error_message = ( - "Error running Dalle3: API Key is missing" - ) - - # Act and assert - with pytest.raises(ValueError) as excinfo: - dalle3.create_variations(img_url) - - assert str(excinfo.value) == expected_error_message - - -def test_dalle3_call_with_retry_max_retries_exceeded( - dalle3, mock_openai_client -): - # Arrange - task = "A painting of a dog" - - # Simulate max retries exceeded - mock_openai_client.images.generate.side_effect = OpenAIError( - "Temporary error", - http_status=500, - error="Internal Server Error", - ) - - # Act and assert - with pytest.raises(OpenAIError) as excinfo: - dalle3(task) - - assert "Retry limit exceeded" in str(excinfo.value) - - -def test_dalle3_create_variations_with_retry_max_retries_exceeded( - dalle3, mock_openai_client -): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - - # Simulate max retries exceeded - mock_openai_client.images.create_variation.side_effect = ( - OpenAIError( - "Temporary error", - http_status=500, - error="Internal Server Error", - ) - ) - - # Act and assert - with pytest.raises(OpenAIError) as excinfo: - dalle3.create_variations(img_url) - - assert "Retry limit exceeded" in str(excinfo.value) - - -def test_dalle3_call_retry_with_success(dalle3, mock_openai_client): - # Arrange - task = "A painting of a dog" - expected_img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - - # Simulate success after a retry - mock_openai_client.images.generate.side_effect = [ - OpenAIError( - "Temporary error", - http_status=500, - error="Internal Server Error", - ), - Mock(data=[Mock(url=expected_img_url)]), - ] - - # Act - img_url = dalle3(task) - - # Assert - assert img_url == expected_img_url - assert mock_openai_client.images.generate.call_count == 2 - - -def test_dalle3_create_variations_retry_with_success( - dalle3, mock_openai_client -): - # Arrange - img_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_01J9J5ZKJZJY9.png" - expected_variation_url = "https://cdn.openai.com/dall-e/encoded/feats/feats_02ABCDE.png" - - # Simulate success after a retry - mock_openai_client.images.create_variation.side_effect = [ - OpenAIError( - "Temporary error", - http_status=500, - error="Internal Server Error", - ), - Mock(data=[Mock(url=expected_variation_url)]), - ] - - # Act - variation_img_url = dalle3.create_variations(img_url) - - # Assert - assert variation_img_url == expected_variation_url - assert mock_openai_client.images.create_variation.call_count == 2