| """Tests for deterministic privacy-safe tracing.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import tempfile |
| import threading |
| import time |
| import unittest |
| from dataclasses import replace |
| from pathlib import Path |
| from unittest.mock import patch |
|
|
| import app |
| from app import config, model_endpoint, service |
| from app import ocr |
| from traces import runtime as trace_runtime |
| from traces.scripts.validate_traces import validate_file |
|
|
|
|
| class TraceTests(unittest.TestCase): |
| def sample_record(self) -> dict: |
| return trace_runtime.build_trace_record( |
| text=( |
| "Ali call +92 300 1234567 CNIC 35202-1234567-1 " |
| "card 4111111111111111 https://private.example/a " |
| "account PK00TEST address House 10 tracking ABC123" |
| ), |
| image_data_url="data:image/png;base64,PRIVATE_IMAGE_BYTES", |
| example_id="", |
| assessment={ |
| "risk_label": "Likely scam", |
| "simple_explanation": "PRIVATE MODEL EXPLANATION", |
| "red_flags": ["PRIVATE FLAG"], |
| "safe_next_steps": ["PRIVATE STEP"], |
| "reply_draft": "PRIVATE REPLY", |
| }, |
| ) |
|
|
| def test_trace_has_no_private_canaries(self) -> None: |
| serialized = json.dumps(self.sample_record()) |
| for canary in ( |
| "Ali", |
| "+92 300 1234567", |
| "35202-1234567-1", |
| "4111111111111111", |
| "private.example", |
| "PK00TEST", |
| "House 10", |
| "ABC123", |
| "PRIVATE_IMAGE_BYTES", |
| "PRIVATE MODEL EXPLANATION", |
| "PRIVATE FLAG", |
| "PRIVATE STEP", |
| "PRIVATE REPLY", |
| ): |
| self.assertNotIn(canary, serialized) |
|
|
| def test_trace_validation_and_performance(self) -> None: |
| started = time.perf_counter() |
| records = [self.sample_record() for _ in range(200)] |
| elapsed_ms = (time.perf_counter() - started) * 1000 / len(records) |
| self.assertLess(elapsed_ms, 10) |
| self.assertFalse(trace_runtime.validate_trace(records[0])) |
|
|
| def test_trace_validation_rejects_extra_or_sensitive_columns(self) -> None: |
| record = self.sample_record() |
| record["notes"] = "PRIVATE RAW MESSAGE" |
| self.assertIn( |
| "Unexpected fields: notes", |
| trace_runtime.validate_trace(record), |
| ) |
|
|
| record = self.sample_record() |
| record["input"] = "text: Visit https://private.example immediately" |
| self.assertIn( |
| "Input contains an unredacted URL.", |
| trace_runtime.validate_trace(record), |
| ) |
|
|
| record = self.sample_record() |
| record["result_summary"] = "PRIVATE MODEL OUTPUT" |
| self.assertIn( |
| "Result summary does not match the deterministic fields.", |
| trace_runtime.validate_trace(record), |
| ) |
|
|
| def test_file_validation_rejects_duplicate_trace_ids(self) -> None: |
| record = self.sample_record() |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / "duplicates.jsonl" |
| path.write_text( |
| json.dumps(record) + "\n" + json.dumps(record) + "\n", |
| encoding="utf-8", |
| ) |
| count, errors = validate_file(path) |
| self.assertEqual(count, 2) |
| self.assertTrue(any("Duplicate trace ID" in error for error in errors)) |
|
|
| def test_trace_uses_simplified_columns(self) -> None: |
| image_record = self.sample_record() |
| self.assertTrue(image_record["input"].startswith("image: ")) |
| self.assertEqual(image_record["input_category"], "unknown") |
| self.assertFalse(image_record["urgency"]) |
| self.assertIn( |
| "does not confirm a new scam type", |
| image_record["result_summary"], |
| ) |
| self.assertIn("Strong scam indicators", image_record["result_summary"]) |
| for removed in ( |
| "pipeline_steps", |
| "cache", |
| "app_commit", |
| "failure", |
| "schema_version", |
| "request_source", |
| "text_byte_bucket", |
| "text_character_bucket", |
| "image_size_bucket", |
| "language_hint", |
| "modal", |
| "exception_text_stored", |
| "identifiers_stored", |
| "input_storage", |
| "raw_image_stored", |
| "raw_input_stored", |
| "raw_model_output_stored", |
| "red_flag_count", |
| "reply_draft_returned", |
| "safe_next_step_count", |
| "signal_account_threat", |
| "signal_challan", |
| "signal_cnic", |
| "signal_courier", |
| "signal_credentials", |
| "signal_link", |
| "signal_otp", |
| "signal_payment", |
| "signal_refund_or_prize", |
| ): |
| self.assertNotIn(removed, image_record) |
|
|
| text_record = trace_runtime.build_trace_record( |
| text="Urgent courier payment required today", |
| image_data_url="", |
| example_id="", |
| assessment=None, |
| ) |
| self.assertTrue(text_record["input"].startswith("text: ")) |
| self.assertIn("courier payment required today", text_record["input"]) |
| self.assertNotIn("Urgent", text_record["input"]) |
| self.assertEqual(text_record["input_category"], "courier") |
| self.assertTrue(text_record["urgency"]) |
| self.assertIn("Known courier pattern", text_record["result_summary"]) |
| self.assertIn( |
| "No completed assessment result was available", |
| text_record["result_summary"], |
| ) |
| self.assertFalse(any( |
| isinstance(value, (dict, list)) |
| for value in text_record.values() |
| )) |
| self.assertEqual(next(iter(text_record)), "trace_id") |
|
|
| def test_image_trace_uses_model_assessment_for_fixed_description(self) -> None: |
| assessment = { |
| "risk_label": "Suspicious", |
| "simple_explanation": ( |
| "The screenshot claims a Pakistan Post parcel delivery failed " |
| "and demands an immediate address update through a short link." |
| ), |
| "red_flags": [ |
| "Unknown sender with no official courier branding.", |
| "Urgent delivery action through a suspicious URL.", |
| ], |
| "safe_next_steps": ["Contact the courier through its official app."], |
| "reply_draft": "", |
| } |
| record = trace_runtime.build_trace_record( |
| text="", |
| image_data_url="data:image/png;base64,PRIVATE_IMAGE_BYTES", |
| example_id="", |
| assessment=assessment, |
| ) |
|
|
| self.assertEqual(record["input_category"], "courier") |
| self.assertTrue(record["urgency"]) |
| self.assertEqual( |
| record["input"], |
| "image: Courier-style content with link, urgency, courier signals", |
| ) |
| self.assertIn("Known courier pattern", record["result_summary"]) |
| serialized = json.dumps(record) |
| self.assertNotIn(assessment["simple_explanation"], serialized) |
| self.assertNotIn(assessment["red_flags"][0], serialized) |
| self.assertNotIn("PRIVATE_IMAGE_BYTES", serialized) |
|
|
| def test_image_trace_uses_result_summary_not_extra_metadata(self) -> None: |
| record = trace_runtime.build_trace_record( |
| text="", |
| image_data_url="data:image/png;base64,PRIVATE_IMAGE_BYTES", |
| example_id="", |
| assessment={ |
| "risk_label": "Suspicious", |
| "simple_explanation": ( |
| "The buyer asks to move the marketplace conversation " |
| "to WhatsApp." |
| ), |
| "red_flags": ["The sender identity is unverified."], |
| "safe_next_steps": ["Verify independently."], |
| "reply_draft": "", |
| "trace_category": "university", |
| "trace_tactics": ["refund_or_prize", "urgency"], |
| }, |
| ) |
|
|
| self.assertEqual(record["input_category"], "marketplace") |
| self.assertEqual( |
| record["scam_tactics"], |
| "off_platform_contact, impersonation", |
| ) |
| self.assertEqual( |
| record["input"], |
| ( |
| "image: Marketplace-style content with off-platform-contact, " |
| "impersonation signals" |
| ), |
| ) |
| self.assertFalse(trace_runtime.validate_trace(record)) |
|
|
| def test_failed_image_trace_is_assessment_unavailable(self) -> None: |
| record = trace_runtime.build_trace_record( |
| text="", |
| image_data_url="data:image/png;base64,PRIVATE_IMAGE_BYTES", |
| example_id="", |
| assessment=None, |
| ) |
|
|
| self.assertEqual(record["input"], "image: Assessment unavailable") |
| self.assertEqual(record["risk_label"], "none") |
| self.assertFalse(trace_runtime.validate_trace(record)) |
|
|
| def test_urdu_image_assessment_maps_to_traffic_challan(self) -> None: |
| record = trace_runtime.build_trace_record( |
| text="", |
| image_data_url="data:image/jpeg;base64,PRIVATE_IMAGE_BYTES", |
| example_id="", |
| assessment={ |
| "risk_label": "Likely scam", |
| "simple_explanation": ( |
| "یہ ٹریفک چالان کی فوری ادائیگی کے لیے مشکوک لنک استعمال کرتا ہے۔" |
| ), |
| "red_flags": ["آج جرمانہ ادا کرنے کا دباؤ"], |
| "safe_next_steps": ["سرکاری ذریعے سے تصدیق کریں۔"], |
| "reply_draft": "", |
| }, |
| ) |
|
|
| self.assertEqual(record["input_category"], "traffic_challan") |
| self.assertTrue(record["urgency"]) |
| self.assertIn("Traffic-challan-style", record["input"]) |
| self.assertIn("link", record["scam_tactics"]) |
| self.assertIn("payment", record["scam_tactics"]) |
|
|
| def test_opt_out_does_not_queue_trace(self) -> None: |
| with patch("app.service.queue_trace") as queue_mock: |
| result = app.analyze_notice("", "", save_trace=False) |
| queue_mock.assert_not_called() |
| self.assertEqual(result["trace"]["status"], "disabled") |
|
|
| def test_cached_trace_does_not_call_model(self) -> None: |
| with patch("app.model_endpoint.call_model") as model_mock, patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ): |
| result = app.analyze_notice(example_id="text-bank") |
| model_mock.assert_not_called() |
| self.assertEqual(result["source"], "cached_local_example") |
|
|
| def test_empty_input_trace_has_no_failure_metadata(self) -> None: |
| with patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ) as queue_mock: |
| result = app.analyze_notice("") |
| self.assertFalse(result["ok"]) |
| self.assertNotIn("failure_category", queue_mock.call_args.kwargs) |
| self.assertNotIn("failure_stage", queue_mock.call_args.kwargs) |
|
|
| def test_missing_credentials_traces_without_model_call(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": False, "label": "missing"}, |
| ), patch("app.model_endpoint.call_model") as model_mock, patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ) as queue_mock: |
| result = app.analyze_notice("test message") |
| self.assertFalse(result["ok"]) |
| model_mock.assert_not_called() |
| self.assertEqual(result["error_code"], "modelConfigurationError") |
| self.assertNotIn("modal_called", queue_mock.call_args.kwargs) |
|
|
| def test_success_uses_existing_model_call_once(self) -> None: |
| assessment = { |
| "risk_label": "Verify first", |
| "simple_explanation": "Check independently.", |
| "red_flags": ["Unverified sender"], |
| "safe_next_steps": ["Use an official channel."], |
| "reply_draft": "Please confirm through your official channel.", |
| } |
|
|
| def fake_call(_text, _image, _language): |
| return assessment |
|
|
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=fake_call, |
| ) as model_mock, patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ) as queue_mock: |
| result = app.analyze_notice("test message") |
| self.assertTrue(result["ok"]) |
| model_mock.assert_called_once() |
| self.assertNotIn("modal_called", queue_mock.call_args.kwargs) |
| self.assertNotIn("retry_count", queue_mock.call_args.kwargs) |
|
|
| def test_trace_failure_does_not_hide_successful_assessment(self) -> None: |
| assessment = { |
| "risk_label": "Verify first", |
| "simple_explanation": "Check independently.", |
| "red_flags": ["Unverified sender"], |
| "safe_next_steps": ["Use an official channel."], |
| "reply_draft": "Please confirm through an official channel.", |
| } |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch("app.model_endpoint.call_model", return_value=assessment), patch( |
| "app.service.queue_trace", |
| side_effect=RuntimeError("trace publisher failed"), |
| ): |
| result = app.analyze_notice("test message", save_trace=True) |
|
|
| self.assertTrue(result["ok"]) |
| self.assertEqual(result["assessment"], assessment) |
| self.assertEqual(result["trace"]["status"], "failed") |
|
|
| def test_timeout_is_sanitized(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=model_endpoint.ModelRuntimeError( |
| "The local GGUF model could not be loaded." |
| ), |
| ), patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ) as queue_mock: |
| result = app.analyze_notice("test message") |
| self.assertFalse(result["ok"]) |
| self.assertEqual(result["error_code"], "modelUnavailableError") |
| self.assertNotIn("failure_category", queue_mock.call_args.kwargs) |
|
|
| def test_http_failure_is_sanitized(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=model_endpoint.ModelRuntimeError("private runtime failure"), |
| ), patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ) as queue_mock: |
| result = app.analyze_notice("test message") |
| self.assertFalse(result["ok"]) |
| self.assertEqual(result["error_code"], "modelUnavailableError") |
| self.assertNotIn("private", json.dumps(queue_mock.call_args.kwargs)) |
|
|
| def test_malformed_output_is_sanitized(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=ValueError("PRIVATE RAW OUTPUT"), |
| ), patch( |
| "app.service.queue_trace", |
| return_value=("trace-id", "queued"), |
| ) as queue_mock: |
| result = app.analyze_notice("test message") |
| self.assertFalse(result["ok"]) |
| self.assertEqual(result["error_code"], "modelInvalidError") |
| self.assertNotIn("PRIVATE RAW OUTPUT", json.dumps(queue_mock.call_args.kwargs)) |
|
|
| def test_image_without_notice_text_returns_input_warning(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=model_endpoint.NoticeImageInputError( |
| "No readable notice text was found in the screenshot." |
| ), |
| ): |
| result = app.analyze_notice( |
| image_data_url="data:image/png;base64,AAAA", |
| save_trace=False, |
| ) |
|
|
| self.assertFalse(result["ok"]) |
| self.assertTrue(result["warning"]) |
| self.assertEqual(result["error_code"], "noticeImageRequiredWarning") |
| self.assertTrue(result["status"]["connected"]) |
|
|
| def test_image_invalid_generation_returns_input_warning(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=RuntimeError("sanitized ZeroGPU worker failure"), |
| ): |
| result = app.analyze_notice( |
| image_data_url="data:image/png;base64,AAAA", |
| save_trace=False, |
| ) |
|
|
| self.assertFalse(result["ok"]) |
| self.assertTrue(result["warning"]) |
| self.assertEqual(result["error_code"], "noticeImageRequiredWarning") |
| self.assertTrue(result["status"]["connected"]) |
|
|
| def test_text_invalid_generation_remains_model_error(self) -> None: |
| with patch( |
| "app.model_endpoint.model_status", |
| return_value={"connected": True, "label": "ready"}, |
| ), patch( |
| "app.model_endpoint.call_model", |
| side_effect=RuntimeError( |
| "The local model returned an invalid response." |
| ), |
| ): |
| result = app.analyze_notice( |
| text="Please pay this bill.", |
| save_trace=False, |
| ) |
|
|
| self.assertFalse(result["ok"]) |
| self.assertEqual(result["error_code"], "modelInvalidError") |
| self.assertNotIn("warning", result) |
|
|
| def test_normalization_failure_uses_normalize_stage(self) -> None: |
| telemetry: dict = {} |
| with self.assertRaises(ValueError): |
| app.parse_model_json('{"risk_label":"invalid"}', telemetry) |
| self.assertTrue(telemetry["parse_completed"]) |
| self.assertNotIn("normalize_completed", telemetry) |
|
|
| def test_local_completion_uses_structured_output(self) -> None: |
| valid = { |
| "risk_label": "Verify first", |
| "simple_explanation": "آزاد ذریعے سے تصدیق کریں۔", |
| "red_flags": ["بھیجنے والے کی تصدیق نہیں ہوئی۔"], |
| "safe_next_steps": ["سرکاری ذریعے سے رابطہ کریں۔"], |
| "reply_draft": "براہ کرم سرکاری ذریعے سے تصدیق کریں۔", |
| } |
|
|
| class LocalModel: |
| def __init__(self): |
| self.calls = 0 |
| self.max_tokens: list[int] = [] |
| self.response_formats: list[dict] = [] |
|
|
| def create_chat_completion(self, **kwargs): |
| self.calls += 1 |
| self.max_tokens.append(kwargs["max_tokens"]) |
| self.response_formats.append(kwargs["response_format"]) |
| return { |
| "choices": [ |
| { |
| "message": { |
| "content": ( |
| "<think>private reasoning</think>\n" |
| + json.dumps(valid) |
| ) |
| } |
| } |
| ] |
| } |
|
|
| local_model = LocalModel() |
| result = model_endpoint._run_completion(local_model, "test", "ur") |
| self.assertEqual(result["risk_label"], "Verify first") |
| self.assertEqual(local_model.calls, 1) |
| self.assertEqual(local_model.max_tokens, [1200]) |
| self.assertEqual(local_model.response_formats[0]["type"], "json_object") |
| self.assertNotIn( |
| '"type": "object"', |
| model_endpoint._messages("test", "en")[1]["content"], |
| ) |
|
|
| def test_ocr_extracts_paragraph_text(self) -> None: |
| import torch |
|
|
| image_data = ( |
| "data:image/png;base64," |
| "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4nGNg" |
| "YAAAAAMAASsJTYQAAAAASUVORK5CYII=" |
| ) |
| parameter = unittest.mock.Mock(device=torch.device("cpu"), dtype=torch.float32) |
| model = unittest.mock.Mock() |
| model.parameters.side_effect = lambda: iter([parameter]) |
| model.generate.return_value = torch.tensor([[1, 2, 3]]) |
| processor = unittest.mock.Mock() |
| processor.return_value = {"pixel_values": torch.ones(1)} |
| processor.batch_decode.return_value = ["model output"] |
| postprocessing = unittest.mock.Mock() |
| postprocessing.extract_classes_bboxes.return_value = ( |
| ["paragraph", "paragraph"], |
| [None, None], |
| ["PAKISTAN POST", "Pay Rs. 85 now"], |
| ) |
| postprocessing.postprocess_text.side_effect = ( |
| lambda text, **_kwargs: text |
| ) |
|
|
| with patch( |
| "app.ocr._get_model", |
| return_value=(model, processor, unittest.mock.Mock()), |
| ), patch( |
| "app.ocr._load_postprocessing", |
| return_value=postprocessing, |
| ): |
| text = ocr.extract_text(image_data) |
|
|
| self.assertEqual(text, "PAKISTAN POST\n\nPay Rs. 85 now") |
| model.generate.assert_called_once() |
|
|
| def test_ocr_readability_rejects_parser_markup_without_notice_text(self) -> None: |
| self.assertFalse( |
| ocr._has_readable_text( |
| "<picture><x_10><y_20><x_900><y_700></picture>" |
| ) |
| ) |
| self.assertTrue(ocr._has_readable_text("Pay Rs. 85 now")) |
| self.assertTrue(ocr._has_readable_text("آپ کا بل 500 روپے ہے")) |
|
|
| def test_ocr_rejects_picture_region_with_generated_description(self) -> None: |
| import torch |
|
|
| image_data = ( |
| "data:image/png;base64," |
| "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4nGNg" |
| "YAAAAAMAASsJTYQAAAAASUVORK5CYII=" |
| ) |
| parameter = unittest.mock.Mock(device=torch.device("cpu"), dtype=torch.float32) |
| model = unittest.mock.Mock() |
| model.parameters.side_effect = lambda: iter([parameter]) |
| model.generate.return_value = torch.tensor([[1, 2, 3]]) |
| processor = unittest.mock.Mock() |
| processor.return_value = {"pixel_values": torch.ones(1)} |
| processor.batch_decode.return_value = ["generated picture output"] |
| postprocessing = unittest.mock.Mock() |
| postprocessing.extract_classes_bboxes.return_value = ( |
| ["Picture"], |
| [None], |
| ["A portrait of a woman"], |
| ) |
|
|
| with patch( |
| "app.ocr._get_model", |
| return_value=(model, processor, unittest.mock.Mock()), |
| ), patch( |
| "app.ocr._load_postprocessing", |
| return_value=postprocessing, |
| ): |
| with self.assertRaises(ocr.NoReadableTextError): |
| ocr.extract_text(image_data) |
|
|
| postprocessing.postprocess_text.assert_not_called() |
|
|
| def test_ocr_keeps_text_regions_next_to_a_picture(self) -> None: |
| self.assertFalse(ocr._is_text_class("Picture")) |
| self.assertFalse(ocr._is_text_class("page-image")) |
| self.assertTrue(ocr._is_text_class("Paragraph")) |
| self.assertTrue(ocr._is_text_class("Table")) |
|
|
| def test_no_text_ocr_error_becomes_notice_image_input_error(self) -> None: |
| with patch( |
| "app.model_endpoint.extract_text", |
| side_effect=ocr.NoReadableTextError( |
| "No readable notice text was found in the screenshot." |
| ), |
| ): |
| with self.assertRaises(model_endpoint.NoticeImageInputError): |
| model_endpoint.call_model( |
| "", |
| "data:image/png;base64,AAAA", |
| ) |
|
|
| def test_image_ocr_text_is_passed_to_minicpm(self) -> None: |
| config = model_endpoint.model_config() |
| fake_model = object() |
| assessment = { |
| "risk_label": "Likely scam", |
| "simple_explanation": "Urgent payment request.", |
| "red_flags": ["Untrusted payment request"], |
| "safe_next_steps": ["Verify through an official channel."], |
| "reply_draft": "", |
| } |
| with patch( |
| "app.model_endpoint.extract_text", |
| return_value="Pay Rs. 85 immediately", |
| ), patch( |
| "app.model_endpoint.model_config", |
| return_value=replace(config, keep_loaded=False), |
| ), patch( |
| "app.model_endpoint._build_model", |
| return_value=fake_model, |
| ), patch( |
| "app.model_endpoint._run_completion", |
| return_value=assessment, |
| ) as completion_mock: |
| result = model_endpoint.call_model( |
| "", |
| ( |
| "data:image/png;base64," |
| "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4" |
| "nGNgYAAAAAMAASsJTYQAAAAASUVORK5CYII=" |
| ), |
| ) |
| self.assertEqual(result, assessment) |
| self.assertEqual( |
| completion_mock.call_args.args[1], |
| "Pay Rs. 85 immediately", |
| ) |
|
|
| def test_space_uses_transformers_instead_of_llama_cpp(self) -> None: |
| config = model_endpoint.model_config() |
| assessment = { |
| "risk_label": "Verify first", |
| "simple_explanation": "Confirm the sender independently.", |
| "red_flags": ["Unverified sender"], |
| "safe_next_steps": ["Use an official contact channel."], |
| "reply_draft": "I will verify this independently.", |
| } |
| with patch.dict("os.environ", {"SPACE_ID": "noticecheck"}), patch( |
| "app.model_endpoint.model_config", |
| return_value=config, |
| ), patch( |
| "app.model_endpoint._run_transformers_completion", |
| return_value=assessment, |
| ) as transformers_mock, patch( |
| "app.model_endpoint._build_model", |
| ) as llama_mock: |
| result = model_endpoint.call_model("Please pay this fee now.") |
|
|
| self.assertEqual(result, assessment) |
| transformers_mock.assert_called_once_with( |
| "Please pay this fee now.", |
| "en", |
| ) |
| llama_mock.assert_not_called() |
|
|
| def test_explicit_local_transformers_runtime_avoids_llama_cpp(self) -> None: |
| assessment = { |
| "risk_label": "Verify first", |
| "simple_explanation": "Confirm the sender independently.", |
| "red_flags": ["Unverified sender"], |
| "safe_next_steps": ["Use an official contact channel."], |
| "reply_draft": "I will verify this independently.", |
| } |
| with patch.dict( |
| "os.environ", |
| {"MODEL_RUNTIME": "transformers"}, |
| clear=False, |
| ), patch( |
| "app.model_endpoint._run_transformers_completion", |
| return_value=assessment, |
| ) as transformers_mock, patch( |
| "app.model_endpoint._build_model", |
| ) as llama_mock: |
| result = model_endpoint.call_model("Please pay this fee now.") |
|
|
| self.assertEqual(result, assessment) |
| transformers_mock.assert_called_once_with( |
| "Please pay this fee now.", |
| "en", |
| ) |
| llama_mock.assert_not_called() |
|
|
| def test_invalid_model_runtime_is_rejected(self) -> None: |
| with patch.dict( |
| "os.environ", |
| {"MODEL_RUNTIME": "invalid"}, |
| clear=False, |
| ): |
| with self.assertRaisesRegex(ValueError, "MODEL_RUNTIME"): |
| config.model_runtime() |
|
|
| def test_transformers_completion_uses_model_card_generation_flow(self) -> None: |
| assessment = { |
| "risk_label": "Likely scam", |
| "simple_explanation": "The message requests an OTP.", |
| "red_flags": ["OTP request"], |
| "safe_next_steps": ["Do not share the OTP."], |
| "reply_draft": "", |
| } |
|
|
| class Encoded(dict): |
| def to(self, device): |
| self["device"] = device |
| return self |
|
|
| class InputIds: |
| shape = (1, 12) |
|
|
| encoded = Encoded(input_ids=InputIds()) |
| tokenizer = unittest.mock.Mock() |
| tokenizer.eos_token_id = 2 |
| tokenizer.apply_chat_template.return_value = encoded |
| tokenizer.decode.return_value = json.dumps(assessment) |
| model = unittest.mock.Mock() |
| model.device = "cuda:0" |
| model.generate.return_value = [[0] * 13] |
|
|
| with patch( |
| "app.model_endpoint._get_transformers_model", |
| return_value=(tokenizer, model), |
| ): |
| result = model_endpoint._run_transformers_completion( |
| "Share your OTP now.", |
| "en", |
| ) |
|
|
| self.assertEqual(result, assessment) |
| self.assertEqual(encoded["device"], "cuda:0") |
| tokenizer.apply_chat_template.assert_called_once() |
| model.generate.assert_called_once() |
|
|
| def test_transformers_completion_repairs_invalid_json_once(self) -> None: |
| assessment = { |
| "risk_label": "Suspicious", |
| "simple_explanation": "The sender is not verified.", |
| "red_flags": ["Unverified sender"], |
| "safe_next_steps": ["Confirm through an official channel."], |
| "reply_draft": "I will verify this independently.", |
| } |
|
|
| class Encoded(dict): |
| def to(self, device): |
| return self |
|
|
| class InputIds: |
| shape = (1, 8) |
|
|
| tokenizer = unittest.mock.Mock() |
| tokenizer.eos_token_id = 2 |
| tokenizer.apply_chat_template.return_value = Encoded( |
| input_ids=InputIds() |
| ) |
| tokenizer.decode.side_effect = [ |
| "This response is not JSON.", |
| json.dumps(assessment), |
| ] |
| model = unittest.mock.Mock() |
| model.device = "cuda:0" |
| model.generate.return_value = [[0] * 9] |
|
|
| with patch( |
| "app.model_endpoint._get_transformers_model", |
| return_value=(tokenizer, model), |
| ): |
| result = model_endpoint._run_transformers_completion( |
| "Confirm this sender.", |
| "en", |
| ) |
|
|
| self.assertEqual(result, assessment) |
| self.assertEqual(model.generate.call_count, 2) |
|
|
| def test_urdu_script_ocr_is_assessed_with_english_output(self) -> None: |
| assessment = { |
| "risk_label": "Verify first", |
| "simple_explanation": "Confirm the sender independently.", |
| "red_flags": ["Unverified sender"], |
| "safe_next_steps": ["Use an official contact channel."], |
| "reply_draft": "I will verify this independently.", |
| } |
| with patch.dict( |
| "os.environ", |
| {"MODEL_RUNTIME": "transformers"}, |
| clear=False, |
| ), patch( |
| "app.model_endpoint.extract_text", |
| return_value="آپ کا اکاؤنٹ بند ہو جائے گا", |
| ), patch( |
| "app.model_endpoint._run_transformers_completion", |
| return_value=assessment, |
| ) as completion_mock: |
| result = model_endpoint.call_model( |
| "", |
| ( |
| "data:image/png;base64," |
| "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lE" |
| "QVR4nGNgYAAAAAMAASsJTYQAAAAASUVORK5CYII=" |
| ), |
| ) |
|
|
| self.assertEqual(result, assessment) |
| completion_mock.assert_called_once_with( |
| "آپ کا اکاؤنٹ بند ہو جائے گا", |
| "en", |
| ) |
|
|
| def test_publisher_persists_batch(self) -> None: |
| publisher = trace_runtime.TracePublisher() |
| with tempfile.TemporaryDirectory() as directory, patch.object( |
| trace_runtime, |
| "PENDING_DIR", |
| Path(directory), |
| ): |
| records = [self.sample_record() for _ in range(20)] |
| publisher._persist_batch(records) |
| paths = list(Path(directory).glob("*.jsonl")) |
| self.assertEqual(len(paths), 1) |
| self.assertEqual(len(paths[0].read_text().splitlines()), 20) |
|
|
| def test_concurrent_enqueue_and_queue_limit(self) -> None: |
| with patch.object(trace_runtime, "MAX_QUEUE_SIZE", 20): |
| publisher = trace_runtime.TracePublisher() |
| with patch.object(publisher, "_ensure_worker"): |
| threads = [ |
| threading.Thread(target=publisher.enqueue, args=(self.sample_record(),)) |
| for _ in range(20) |
| ] |
| for thread in threads: |
| thread.start() |
| for thread in threads: |
| thread.join() |
| self.assertEqual(publisher.status()["queued"], 20) |
| self.assertEqual(publisher.queue.qsize(), 20) |
| self.assertEqual(publisher.enqueue(self.sample_record()), "dropped") |
| self.assertEqual(publisher.status()["dropped"], 1) |
|
|
| def test_hub_failure_keeps_pending_shard(self) -> None: |
| publisher = trace_runtime.TracePublisher() |
| with tempfile.TemporaryDirectory() as directory, patch.object( |
| trace_runtime, |
| "PENDING_DIR", |
| Path(directory), |
| ), patch.dict("os.environ", {"HF_TOKEN": "test-token"}), patch( |
| "huggingface_hub.HfApi.upload_file", |
| side_effect=RuntimeError("offline"), |
| ), patch("traces.runtime.time.sleep"): |
| publisher._persist_batch([self.sample_record()]) |
| publisher._upload_pending() |
| self.assertEqual(len(list(Path(directory).glob("*.jsonl"))), 1) |
| self.assertEqual(publisher.status()["upload_failures"], 1) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
