File size: 2,223 Bytes
84e50e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import unittest
import os
import json
from PIL import Image
from pii_image_processing import process_image_api, PIIExtractorFactory
class DummyExtractor:
def __init__(self, model=None):
pass
def extract_pii(self, image_input):
sample = {
"piis": [
{
"name": "TestPII",
"coordinates": {"x1": 10, "y1": 10, "x2": 50, "y2": 50},
"probable_regulations": ["GDPR"]
}
],
"containing_text": "TestPII"
}
return json.dumps(sample)
class TestProcessImageApi(unittest.TestCase):
@classmethod
def setUpClass(cls):
# Monkey-patch factory to use dummy extractor
PIIExtractorFactory.create_extractor = staticmethod(lambda provider, **kwargs: DummyExtractor(**kwargs))
os.makedirs("tmp", exist_ok=True)
cls.test_image = "tmp/dummy_test.jpg"
Image.new("RGB", (100, 100), (128, 128, 128)).save(cls.test_image)
def test_blur_strategy(self):
out = "tmp/output_blur.jpg"
result = process_image_api(
self.test_image,
strategy_name="blur",
blur_amount=2,
output_path=out
)
self.assertTrue(result.get("success"))
self.assertTrue(os.path.exists(out))
self.assertEqual(len(result["data"]), 1)
def test_single_color_strategy(self):
out = "tmp/output_color.jpg"
result = process_image_api(
self.test_image,
strategy_name="single_color",
color=(255,0,0),
output_path=out
)
self.assertTrue(result.get("success"))
self.assertTrue(os.path.exists(out))
self.assertEqual(len(result["data"]), 1)
def test_regulation_map(self):
out = "tmp/output_reg.jpg"
reg_map = {"GDPR": "single_color"}
result = process_image_api(
self.test_image,
regulation_map=reg_map,
output_path=out
)
self.assertTrue(result.get("success"))
self.assertTrue(os.path.exists(out))
self.assertEqual(len(result["data"]), 1)
if __name__ == "__main__":
unittest.main()
|