tests: Add cases for /redact* endpoints

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

tests/integration/assets/public_key.pem

@@ -0,0 +1,3 @@
+-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA3ITkTP8Tm/5FygcwY2EQ7LgVsuCF0OH7psUqvlXnOPNCfX86CobHBiSFjG9o5ZeajPtTXaf1thUodgpJZVZSqpVTXwGKo8r0COMO87IcwYigkZZgG/WmZgoZART+AA0+JvjFGxflJAxSv7puGlf82E+u5Wz2psLBSDO5qrnmaDZTvPh5eX84cocahVVI7X09/kI+sZiKauM69yoy1bdx16YIIeNm0M9qqS3tTrjouQiJfZ8jUKSZ44Na/81LMVw5O46+5GvwD+OsR43kQ0TexMwgtHxQQsiXLWHCDNy2ZzkzukDYRwA3V2lwVjtQN0WjxHg24BTBDBM+v7iQ7cbweQIDAQAB
+-----END PUBLIC KEY-----

tests/integration/test_api.py

@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 
 import pytest
@@ -26,6 +27,7 @@
 
 TEST_ASSETS_DIR = Path("tests/integration/assets")
 MULTI_TEXT_FILE_PATH = TEST_ASSETS_DIR / "multi_text_file.json"
+PUBLIC_KEY_PEM_PATH = TEST_ASSETS_DIR / "public_key.pem"
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -283,3 +285,69 @@ def test_process_bulk_file(client: TestClient, config: Config, test_model_servic
         assert annotation["label_name"] == "Loss Of Kidney Function"
 
     verify_results_match_api_info(client, task, res)
+
+
+def test_redact(client: TestClient, config: Config, test_model_service_ip: str):
+    response = client.post(
+        f"/models/{test_model_service_ip}/redact",
+        data="Patient diagnosed with kidney failure",
+        headers={"Content-Type": "text/plain"},
+    )
+    response_json = validate_api_response(response, expected_status_code=200, return_json=True)
+
+    tm: TaskManager = config.task_manager
+    verify_task_submitted_successfully(response_json["uuid"], tm)
+
+    task = wait_for_task_completion(response_json["uuid"], tm, expected_status=Status.SUCCEEDED)
+
+    key = f"{task.uuid}_payload.txt"
+    expected_payload = b"Patient diagnosed with kidney failure"
+    verify_task_payload_in_object_store(key, expected_payload, config.task_object_store_manager)
+
+    verify_queue_is_empty(config.queue_manager)
+
+    res, parsed = download_result_object(task.result, config.results_object_store_manager, "text")
+
+    assert parsed == "Patient diagnosed with [Loss Of Kidney Function]"
+
+    verify_results_match_api_info(client, task, res)
+
+
+def test_redact_with_encryption(client: TestClient, config: Config, test_model_service_ip: str):
+    with open(PUBLIC_KEY_PEM_PATH) as f:
+        payload = {
+            "text": "Patient diagnosed with kidney failure",
+            "public_key_pem": f.read(),
+        }
+    response = client.post(
+        f"/models/{test_model_service_ip}/redact_with_encryption",
+        json=payload,
+    )
+    response_json = validate_api_response(response, expected_status_code=200, return_json=True)
+
+    tm: TaskManager = config.task_manager
+    verify_task_submitted_successfully(response_json["uuid"], tm)
+
+    task = wait_for_task_completion(response_json["uuid"], tm, expected_status=Status.SUCCEEDED)
+
+    key = f"{task.uuid}_payload.json"
+    expected_payload = json.dumps(payload).encode()
+    verify_task_payload_in_object_store(key, expected_payload, config.task_object_store_manager)
+
+    verify_queue_is_empty(config.queue_manager)
+
+    res, parsed = download_result_object(task.result, config.results_object_store_manager)
+
+    encrypted_label = "[REDACTED_0]"
+    assert "redacted_text" in parsed
+    assert parsed["redacted_text"] == f"Patient diagnosed with {encrypted_label}"
+    assert "encryptions" in parsed
+    assert len(parsed["encryptions"]) == 1
+
+    assert "label" in parsed["encryptions"][0]
+    assert parsed["encryptions"][0]["label"] == encrypted_label
+    assert "encryption" in parsed["encryptions"][0]
+    assert isinstance(parsed["encryptions"][0]["encryption"], str)
+    assert len(parsed["encryptions"][0]["encryption"]) > 0
+
+    verify_results_match_api_info(client, task, res)

tests/integration/utils.py

@@ -315,6 +315,8 @@ def download_result_object(
             result_json = json.loads(result.decode("utf-8"))
         elif format == "jsonl":
             result_json = [json.loads(line) for line in result.decode("utf-8").split("\n") if line]
+        elif format == "text":
+            result_json = result.decode("utf-8")
         else:
             pytest.fail(f"Unsupported format: {format}")
     except json.JSONDecodeError as e: