-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fixes #57: add unittest for image validation
- Loading branch information
Maxence Guindon
committed
Mar 20, 2024
1 parent
34e294e
commit 7f81fae
Showing
3 changed files
with
139 additions
and
80 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,76 +1,135 @@ | ||
| import unittest | ||
| import base64 | ||
| import json | ||
| import asyncio | ||
|
|
||
| from app import app | ||
| from io import BytesIO | ||
| from PIL import Image | ||
| from unittest import TestCase, main | ||
| import requests | ||
|
|
||
| """ | ||
| In order for the tests to run, the server must be running. | ||
| TO DO - Create a mock server to run the tests without the need for the server to be running. | ||
| TO DO - Implement test_image for every other checks (size, format, resizable) | ||
| TO DO - Implement test_image for every type of image (PNG, JPEG, GIF, BMP, TIFF, WEBP, SVG) | ||
| TO DO - | ||
| """ | ||
|
|
||
| class test_image_validation(TestCase): | ||
| # V1 with server running | ||
| def test_real_image_validation(self): | ||
| image = Image.new('RGB', (150, 150), 'blue') | ||
| from unittest.mock import patch, Mock | ||
|
|
||
| # Save the image to a byte array | ||
| img_byte_array = BytesIO() | ||
|
|
||
| image_header = "data:image/PNG;base64," | ||
| class TestImageValidation(unittest.TestCase): | ||
| def setUp(self): | ||
| self.test_client = app.test_client() | ||
|
|
||
| image.save(img_byte_array, 'PNG') | ||
| self.img_byte_array = BytesIO() | ||
| image = Image.new('RGB', (150, 150), 'blue') | ||
| self.image_header = "data:image/PNG;base64," | ||
| image.save(self.img_byte_array, 'PNG') | ||
|
|
||
| data = base64.b64encode(img_byte_array.getvalue()).decode('utf-8') | ||
| def test_real_image_validation(self): | ||
| data = base64.b64encode(self.img_byte_array.getvalue()) | ||
| data = data.decode('utf-8') | ||
|
|
||
| response = requests.post( | ||
| url="http://0.0.0.0:8080/image-validation", | ||
| data= str.encode(json.dumps({'image': image_header + data})), | ||
| headers={ | ||
| response = asyncio.run( | ||
| self.test_client.post( | ||
| '/image-validation', | ||
| headers={ | ||
| "Content-Type": "application/json", | ||
| "Access-Control-Allow-Origin": "*", | ||
| } | ||
| ) | ||
| data = json.loads(response.content) | ||
| }, | ||
| data= str.encode(json.dumps({'image': self.image_header + data})), | ||
| )) | ||
|
|
||
| if isinstance(data[1], str): | ||
| self.assertEqual(response.status_code, 200) | ||
| self.assertEqual(data[0], True) | ||
| else: | ||
| self.assertEqual(response.status_code, 200) | ||
| data = json.loads(asyncio.run(response.get_data())) | ||
|
|
||
| self.assertEqual(response.status_code, 200) | ||
| self.assertIsInstance(data[0], str) | ||
|
|
||
| # v2 with server not running | ||
| def test_invalid_header_image_validation(self): | ||
| image = Image.new('RGB', (150, 150), 'blue') | ||
| data = base64.b64encode(self.img_byte_array.getvalue()).decode('utf-8') | ||
|
|
||
| response = asyncio.run( | ||
| self.test_client.post( | ||
| '/image-validation', | ||
| headers={ | ||
| "Content-Type": "application/json", | ||
| "Access-Control-Allow-Origin": "*", | ||
| }, | ||
| data= str.encode(json.dumps({'image':"data:image/," + data})), | ||
| )) | ||
|
|
||
| data = json.loads(asyncio.run(response.get_data())) | ||
|
|
||
| self.assertEqual(response.status_code, 400) | ||
| self.assertEqual(data[0], 'Invalid file header') | ||
|
|
||
| @patch("PIL.Image.open") | ||
| def test_invalid_extension(self, mock_open): | ||
|
|
||
| mock_image = Mock() | ||
| mock_image.format = "md" | ||
|
|
||
| mock_open.return_value = mock_image | ||
|
|
||
| # Save the image to a byte array | ||
| img_byte_array = BytesIO() | ||
| data = base64.b64encode(self.img_byte_array.getvalue()).decode('utf-8') | ||
|
|
||
| response = asyncio.run( | ||
| self.test_client.post( | ||
| '/image-validation', | ||
| headers={ | ||
| "Content-Type": "application/json", | ||
| "Access-Control-Allow-Origin": "*", | ||
| }, | ||
| data= str.encode(json.dumps({'image': self.image_header + data})), | ||
| )) | ||
|
|
||
| data = json.loads(asyncio.run(response.get_data())) | ||
|
|
||
| self.assertEqual(response.status_code, 400) | ||
| self.assertEqual(data[0], 'Invalid file extension') | ||
|
|
||
| @patch("PIL.Image.open") | ||
| def test_invalid_size(self, mock_open): | ||
| mock_image = Mock() | ||
| mock_image.size = [2000, 2000] | ||
| mock_image.format = "PNG" | ||
|
|
||
| mock_open.return_value = mock_image | ||
|
|
||
| data = base64.b64encode(self.img_byte_array.getvalue()).decode('utf-8') | ||
|
|
||
| response = asyncio.run( | ||
| self.test_client.post( | ||
| '/image-validation', | ||
| headers={ | ||
| "Content-Type": "application/json", | ||
| "Access-Control-Allow-Origin": "*", | ||
| }, | ||
| data= str.encode(json.dumps({'image': self.image_header + data})), | ||
| )) | ||
|
|
||
| data = json.loads(asyncio.run(response.get_data())) | ||
|
|
||
| self.assertEqual(response.status_code, 400) | ||
| self.assertEqual(data[0], 'Invalid file size') | ||
|
|
||
| image_header = "data:image/," | ||
| @patch("PIL.Image.open") | ||
| def test_rezisable_error(self, mock_open): | ||
| mock_image = Mock() | ||
| mock_image.size = [1080, 1080] | ||
| mock_image.format = "PNG" | ||
| mock_image.thumbnail.side_effect = IOError("error can't resize") | ||
|
|
||
| image.save(img_byte_array, 'PNG') | ||
| mock_open.return_value = mock_image | ||
|
|
||
| data = base64.b64encode(img_byte_array.getvalue()).decode('utf-8') | ||
| data = base64.b64encode(self.img_byte_array.getvalue()).decode('utf-8') | ||
|
|
||
| response = requests.post( | ||
| url="http://0.0.0.0:8080/image-validation", | ||
| data= str.encode(json.dumps({'image': image_header + data})), | ||
| headers={ | ||
| response = asyncio.run( | ||
| self.test_client.post( | ||
| '/image-validation', | ||
| headers={ | ||
| "Content-Type": "application/json", | ||
| "Access-Control-Allow-Origin": "*", | ||
| } | ||
| ) | ||
| }, | ||
| data= str.encode(json.dumps({'image': self.image_header + data})), | ||
| )) | ||
|
|
||
| data = json.loads(response.content) | ||
| data = json.loads(asyncio.run(response.get_data())) | ||
|
|
||
| self.assertEqual(response.status_code, 400) | ||
| self.assertEqual(data[0], False) | ||
| self.assertEqual(data[1], 'Invalid file header') | ||
| self.assertEqual(data[0], 'Invalid file not resizable') | ||
|
|
||
| if __name__ == '__main__': | ||
| main() | ||
| unittest.main() |