Update and add more unit tests

unit_tests.py

@@ -1,86 +1,214 @@
 import unittest
 from unittest.mock import patch, MagicMock
+import os
 from PIL import Image
-import base64
-import numpy as np
 from io import BytesIO
+import numpy as np
 from handler import EndpointHandler
 
 class TestEndpointHandler(unittest.TestCase):
-
     @patch('handler.RealESRGANer')
-    def setUp(self, mock_RealESRGANer):
-        self.handler = EndpointHandler(path=".")
+    @patch('handler.boto3')
+    def setUp(self, mock_boto3, mock_RealESRGANer):
+        """Set up test environment before each test"""
+        # Set required environment variables
+        os.environ['TILING_SIZE'] = '0'
+        os.environ['AWS_ACCESS_KEY_ID'] = 'test_key'
+        os.environ['AWS_SECRET_ACCESS_KEY'] = 'test_secret'
+        os.environ['S3_BUCKET_NAME'] = 'test-bucket'
+        
+        self.handler = EndpointHandler()
         self.mock_model = mock_RealESRGANer.return_value
+        self.mock_s3 = mock_boto3.client.return_value
 
-    def create_test_image(self, mode='RGB', size=(100, 100)):
-        image = Image.new(mode, size)
+    def image_to_bytes(self, image):
+        """Helper method to convert PIL Image to bytes"""
         buffered = BytesIO()
         image.save(buffered, format="PNG")
-        return base64.b64encode(buffered.getvalue()).decode()
+        return buffered.getvalue()
 
-    def get_svg_image(self):
-        test_image = "test_data/834989.svg"
-        return test_image
-    
-    def test_float_outscale(self):
-        test_image = self.create_test_image()
-        input_data = {"inputs": {"image": test_image, "outscale": 2.5}}
-        
-        self.mock_model.enhance.return_value = (np.zeros((250, 250, 3), dtype=np.uint8), None)
+    @patch('handler.requests.get')
+    def test_successful_upscale(self, mock_get):
+        """Test successful image upscaling"""
+        # Create test image and mock response
+        test_image = Image.new('RGB', (100, 100))
+        mock_response = MagicMock()
+        mock_response.content = self.image_to_bytes(test_image)
+        mock_get.return_value = mock_response
+        
+        # Mock model output
+        self.mock_model.enhance.return_value = (np.zeros((200, 200, 3), dtype=np.uint8), None)
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 2
+            }
+        }
+        
         result = self.handler(input_data)
         
-        self.assertIn("out_image", result)
+        self.assertIsNotNone(result["image_url"])
+        self.assertIsNotNone(result["image_key"])
         self.assertIsNone(result["error"])
-    
-    def test_outscale_too_small(self):
-        test_image = self.create_test_image()
-        input_data = {"inputs": {"image": test_image, "outscale": 0.5}}
+
+    @patch('handler.requests.get')
+    def test_invalid_outscale(self, mock_get):
+        """Test handling of invalid outscale values"""
+        # Create test image and mock response
+        test_image = Image.new('RGB', (100, 100))
+        mock_response = MagicMock()
+        mock_response.content = self.image_to_bytes(test_image)
+        mock_get.return_value = mock_response
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 0.5  # Too small
+            }
+        }
         
         result = self.handler(input_data)
         
-        self.assertIsNone(result["out_image"])
+        self.assertIsNone(result["image_url"])
+        self.assertIsNone(result["image_key"])
         self.assertIn("Outscale must be between 1 and 10", result["error"])
-    
-    def test_outscale_too_large(self):
-        test_image = self.create_test_image()
-        input_data = {"inputs": {"image": test_image, "outscale": 11}}
+
+    @patch('handler.requests.get')
+    def test_download_failure(self, mock_get):
+        """Test handling of failed image downloads"""
+        mock_get.side_effect = Exception("Download failed")
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 2
+            }
+        }
         
         result = self.handler(input_data)
         
-        self.assertIsNone(result["out_image"])
-        self.assertIn("Outscale must be between 1 and 10", result["error"])
+        self.assertIsNone(result["image_url"])
+        self.assertIsNone(result["image_key"])
+        self.assertIn("Failed to download image", result["error"])
+
+    @patch('handler.requests.get')
+    def test_large_image_no_tiling(self, mock_get):
+        """Test handling of large images when tiling is disabled"""
+        # Create an image larger than max_image_size
+        test_image = Image.new('RGB', (1500, 1500))
+        mock_response = MagicMock()
+        mock_response.content = self.image_to_bytes(test_image)
+        mock_get.return_value = mock_response
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 2
+            }
+        }
         
-    def test_valid_rgb_image(self):
-        test_image = self.create_test_image()
-        input_data = {"inputs": {"image": test_image, "outscale": 2}}
+        result = self.handler(input_data)
+        
+        self.assertIsNone(result["image_url"])
+        self.assertIsNone(result["image_key"])
+        self.assertIn("Image is too large", result["error"])
+
+    @patch('handler.requests.get')
+    def test_s3_upload_failure(self, mock_get):
+        """Test handling of S3 upload failures"""
+        # Create test image and mock response
+        test_image = Image.new('RGB', (100, 100))
+        mock_response = MagicMock()
+        mock_response.content = self.image_to_bytes(test_image)
+        mock_get.return_value = mock_response
         
+        # Mock model output
         self.mock_model.enhance.return_value = (np.zeros((200, 200, 3), dtype=np.uint8), None)
         
+        # Mock S3 upload failure
+        self.mock_s3.upload_fileobj.side_effect = Exception("Upload failed")
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 2
+            }
+        }
+        
         result = self.handler(input_data)
         
-        self.assertIn("out_image", result)
-        self.assertIsNone(result["error"])
-        self.mock_model.enhance.assert_called_once()
+        self.assertIsNone(result["image_url"])
+        self.assertIsNone(result["image_key"])
+        self.assertIn("Failed to upload image to s3", result["error"])
 
-    def test_valid_rgba_image(self):
-        test_image = self.create_test_image(mode='RGBA')
-        input_data = {"inputs": {"image": test_image, "outscale": 2}}
+    def test_missing_image_url(self):
+        """Test handling of missing image URL"""
+        input_data = {
+            "inputs": {
+                "outscale": 2
+            }
+        }
+        
+        result = self.handler(input_data)
         
-        self.mock_model.enhance.return_value = (np.zeros((400, 400, 4), dtype=np.uint8), None)
+        # Check if result contains all required keys
+        self.assertIn("image_url", result)
+        self.assertIn("image_key", result)
+        self.assertIn("error", result)
+        
+        # Check if values are as expected
+        self.assertIsNone(result["image_url"])
+        self.assertIsNone(result["image_key"])
+        self.assertIn("Failed to get inputs", result["error"])
+
+    @patch('handler.requests.get')
+    def test_grayscale_image(self, mock_get):
+        """Test handling of grayscale images"""
+        test_image = Image.new('L', (100, 100))
+        mock_response = MagicMock()
+        mock_response.content = self.image_to_bytes(test_image)
+        mock_get.return_value = mock_response
+        
+        # Mock model output
+        self.mock_model.enhance.return_value = (np.zeros((200, 200), dtype=np.uint8), None)
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 2
+            }
+        }
         
         result = self.handler(input_data)
         
-        self.assertIn("out_image", result)
+        self.assertIsNotNone(result["image_url"])
+        self.assertIsNotNone(result["image_key"])
         self.assertIsNone(result["error"])
 
-    def test_missing_image_key(self):
-        input_data = {"inputs": {}}
+    @patch('handler.requests.get')
+    def test_rgba_image(self, mock_get):
+        """Test handling of RGBA images"""
+        test_image = Image.new('RGBA', (100, 100))
+        mock_response = MagicMock()
+        mock_response.content = self.image_to_bytes(test_image)
+        mock_get.return_value = mock_response
+        
+        # Mock model output
+        self.mock_model.enhance.return_value = (np.zeros((200, 200, 4), dtype=np.uint8), None)
+        
+        input_data = {
+            "inputs": {
+                "image_url": "http://example.com/test.png",
+                "outscale": 2
+            }
+        }
         
         result = self.handler(input_data)
         
-        self.assertIsNone(result["out_image"])
-        self.assertIn("Missing key", result["error"])
+        self.assertIsNotNone(result["image_url"])
+        self.assertIsNotNone(result["image_key"])
+        self.assertIsNone(result["error"])
 
 if __name__ == '__main__':
     unittest.main()