tests/test_image_processor.py

import unittest
from unittest.mock import patch, MagicMock, mock_open
from io import BytesIO
from PIL import Image
from src.image.processor import ImageProcessor
from src.utils.constants import Constants

class TestImageProcessor(unittest.TestCase):
    """
    Test cases for the ImageProcessor class.
    """
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    @patch('src.image.processor.OpenAI')
    def test_init_success(self, mock_openai, mock_getenv, mock_load_dotenv):
        """
        Test successful initialization of ImageProcessor.
        """
        # Setup mocks
        mock_getenv.return_value = "test_key"
        mock_client = MagicMock()
        mock_openai.return_value = mock_client
        
        # Create instance
        processor = ImageProcessor()
        
        # Assertions
        mock_load_dotenv.assert_called_once()
        mock_getenv.assert_called_once_with("OPENAI_API_KEY")
        mock_openai.assert_called_once_with(api_key="test_key")
        self.assertEqual(processor.client, mock_client)
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    def test_init_missing_api_key(self, mock_getenv, mock_load_dotenv):
        """
        Test initialization with missing OpenAI API key.
        """
        # Setup mocks
        mock_getenv.return_value = None
        
        # Assert that ValueError is raised
        with self.assertRaises(ValueError) as context:
            ImageProcessor()
        
        self.assertEqual(str(context.exception), "OPENAI_API_KEY environment variable is not set")
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    @patch('src.image.processor.OpenAI')
    @patch('src.image.processor.Image.open')
    @patch('builtins.open', new_callable=mock_open, read_data=b'test_image_data')
    def test_describe_image_file_path(self, mock_file_open, mock_image_open, mock_openai, 
                                      mock_getenv, mock_load_dotenv):
        """
        Test describe_image method with a file path.
        """
        # Setup mocks
        mock_getenv.return_value = "test_key"
        mock_client = MagicMock()
        mock_openai.return_value = mock_client
        
        mock_pil_image = MagicMock()
        mock_image_open.return_value = mock_pil_image
        
        mock_response = MagicMock()
        mock_choice = MagicMock()
        mock_message = MagicMock()
        mock_message.content = "This is a description of the image."
        mock_choice.message = mock_message
        mock_response.choices = [mock_choice]
        mock_client.chat.completions.create.return_value = mock_response
        
        # Create instance and call method
        processor = ImageProcessor()
        result = processor.describe_image("test_image.jpg")
        
        # Assertions
        mock_file_open.assert_called_once_with("test_image.jpg", "rb")
        mock_image_open.assert_called_once()
        mock_pil_image.thumbnail.assert_called_once()
        mock_client.chat.completions.create.assert_called_once()
        self.assertEqual(result, "This is a description of the image.")
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    @patch('src.image.processor.OpenAI')
    @patch('src.image.processor.Image.open')
    def test_describe_image_bytes(self, mock_image_open, mock_openai, mock_getenv, mock_load_dotenv):
        """
        Test describe_image method with image bytes.
        """
        # Setup mocks
        mock_getenv.return_value = "test_key"
        mock_client = MagicMock()
        mock_openai.return_value = mock_client
        
        mock_pil_image = MagicMock()
        mock_image_open.return_value = mock_pil_image
        
        mock_response = MagicMock()
        mock_choice = MagicMock()
        mock_message = MagicMock()
        mock_message.content = "This is a description of the image."
        mock_choice.message = mock_message
        mock_response.choices = [mock_choice]
        mock_client.chat.completions.create.return_value = mock_response
        
        # Create instance and call method
        processor = ImageProcessor()
        result = processor.describe_image(b'test_image_data')
        
        # Assertions
        mock_image_open.assert_called_once()
        mock_pil_image.thumbnail.assert_called_once()
        mock_client.chat.completions.create.assert_called_once()
        self.assertEqual(result, "This is a description of the image.")
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    @patch('src.image.processor.OpenAI')
    def test_describe_image_none(self, mock_openai, mock_getenv, mock_load_dotenv):
        """
        Test describe_image method with None input.
        """
        # Setup mocks
        mock_getenv.return_value = "test_key"
        
        # Create instance and call method
        processor = ImageProcessor()
        result = processor.describe_image(None)
        
        # Assertions
        self.assertEqual(result, "")
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    @patch('src.image.processor.OpenAI')
    @patch('src.image.processor.Image.open')
    def test_resize_image(self, mock_image_open, mock_openai, mock_getenv, mock_load_dotenv):
        """
        Test resize_image method.
        """
        # Setup mocks
        mock_getenv.return_value = "test_key"
        
        # Create a test image
        test_image = Image.new('RGB', (500, 500), color='red')
        mock_image_open.return_value = test_image
        
        # Create a BytesIO object to capture the saved image
        output = BytesIO()
        
        # Mock the save method to write to our BytesIO object
        original_save = test_image.save
        test_image.save = lambda *args, **kwargs: original_save(output, format='JPEG')
        
        # Create instance and call method
        processor = ImageProcessor()
        result = processor.resize_image(b'test_image_data', 200, 200)
        
        # Assertions
        mock_image_open.assert_called_once()
        self.assertIsNotNone(result)
        
        # Verify the image was resized
        resized_image = Image.open(BytesIO(result))
        self.assertLessEqual(resized_image.width, 200)
        self.assertLessEqual(resized_image.height, 200)
    
    @patch('src.image.processor.load_dotenv')
    @patch('src.image.processor.os.getenv')
    @patch('src.image.processor.OpenAI')
    def test_resize_image_none(self, mock_openai, mock_getenv, mock_load_dotenv):
        """
        Test resize_image method with None input.
        """
        # Setup mocks
        mock_getenv.return_value = "test_key"
        
        # Create instance and call method
        processor = ImageProcessor()
        result = processor.resize_image(None)
        
        # Assertions
        self.assertIsNone(result)

if __name__ == '__main__':
    unittest.main()