add unit tests, create endpoint test notebook

test_handler.ipynb

@@ -0,0 +1,128 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional_tensor.py:5: UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in 0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in torchvision.transforms.v2.functional.\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "from handler import EndpointHandler\n",
+    "import base64\n",
+    "from io import BytesIO\n",
+    "from PIL import Image\n",
+    "import cv2\n",
+    "import random\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# helper decoder\n",
+    "def decode_base64_image(image_string):\n",
+    "  base64_image = base64.b64decode(image_string)\n",
+    "  buffer = BytesIO(base64_image)\n",
+    "  return  Image.open(buffer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# init handler\n",
+    "my_handler = EndpointHandler(path=\".\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "image.size: (1200, 517), image.mode: RGBA, outscale: 10.0\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "output.shape: (5170, 12000, 4)\n",
+      "out_image.size: (12000, 5170)\n",
+      "image.size: (1056, 1068), image.mode: RGB, outscale: 3.0\n",
+      "output.shape: (3204, 3168, 3)\n",
+      "out_image.size: (3168, 3204)\n",
+      "image.size: (1056, 1068), image.mode: L, outscale: 5.49\n",
+      "output.shape: (5863, 5797, 3)\n",
+      "out_image.size: (5797, 5863)\n"
+     ]
+    }
+   ],
+   "source": [
+    "img_dir = \"test_data/\"\n",
+    "img_names = [\"4121783.png\", \"FB_IMG_1725931665635.jpg\", \"FB_IMG_1725931665635_gray.jpg\"]\n",
+    "out_scales = [10, 3, 5.49]\n",
+    "for img_name, outscale in zip(img_names, out_scales):\n",
+    "    image_path = img_dir + img_name\n",
+    "    # create payload\n",
+    "    with open(image_path, \"rb\") as i:\n",
+    "        b64 = base64.b64encode(i.read())\n",
+    "        b64 = b64.decode(\"utf-8\")\n",
+    "        payload = {\n",
+    "            \"inputs\": {\"image\": b64, \n",
+    "                    \"outscale\": outscale\n",
+    "                    }\n",
+    "            }\n",
+    "\n",
+    "\n",
+    "    output_payload = my_handler(payload)\n",
+    "    out_image = decode_base64_image(output_payload[\"out_image\"])\n",
+    "    print(f\"out_image.size: {out_image.size}\")\n",
+    "    out_image.save(f\"test_data/outputs/{img_name.split('.')[0]}_outscale_{outscale}.png\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

unit_tests.py

@@ -0,0 +1,95 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from PIL import Image
+import base64
+import numpy as np
+from io import BytesIO
+from handler import EndpointHandler
+
+class TestEndpointHandler(unittest.TestCase):
+
+    @patch('handler.RealESRGANer')
+    def setUp(self, mock_RealESRGANer):
+        self.handler = EndpointHandler(path=".")
+        self.mock_model = mock_RealESRGANer.return_value
+
+    def create_test_image(self, mode='RGB', size=(100, 100)):
+        image = Image.new(mode, size)
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode()
+
+    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)
+        result = self.handler(input_data)
+        
+        self.assertIn("out_image", result)
+        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}}
+        
+        result = self.handler(input_data)
+        
+        self.assertIsNone(result["out_image"])
+        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}}
+        
+        result = self.handler(input_data)
+        
+        self.assertIsNone(result["out_image"])
+        self.assertIn("Outscale must be between 1 and 10", result["error"])
+        
+    def test_valid_rgb_image(self):
+        test_image = self.create_test_image()
+        input_data = {"inputs": {"image": test_image, "outscale": 2}}
+        
+        self.mock_model.enhance.return_value = (np.zeros((200, 200, 3), dtype=np.uint8), None)
+        
+        result = self.handler(input_data)
+        
+        self.assertIn("out_image", result)
+        self.assertIsNone(result["error"])
+        self.mock_model.enhance.assert_called_once()
+
+    def test_valid_rgba_image(self):
+        test_image = self.create_test_image(mode='RGBA')
+        input_data = {"inputs": {"image": test_image, "outscale": 2}}
+        
+        self.mock_model.enhance.return_value = (np.zeros((400, 400, 4), dtype=np.uint8), None)
+        
+        result = self.handler(input_data)
+        
+        self.assertIn("out_image", result)
+        self.assertIsNone(result["error"])
+
+    def test_image_too_large(self):
+        test_image = self.create_test_image(size=(1500, 1500))
+        input_data = {"inputs": {"image": test_image}}
+        
+        result = self.handler(input_data)
+        
+        self.assertIsNone(result["out_image"])
+        self.assertIn("Image is too large", result["error"])
+
+    def test_missing_image_key(self):
+        input_data = {"inputs": {}}
+        
+        result = self.handler(input_data)
+        
+        self.assertIsNone(result["out_image"])
+        self.assertIn("Missing key", result["error"])
+
+if __name__ == '__main__':
+    unittest.main()