|
|
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): |
|
|
|
|
|
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() |
|
|
|
