first commit

apps/inspector/test_engine.py

@@ -0,0 +1,41 @@
+import io
+import os
+import unittest
+import zipfile
+from engine import cosine, inspect_content, content_kind
+
+def policy(): return {"dlp":{"maxBytes":1048576,"rules":[{"id":"secret","kind":"literal","value":"Secret Roadmap"}],"semantic":{"enabled":False}}}
+
+class InspectionTests(unittest.TestCase):
+    def test_text_and_binary(self):
+        self.assertEqual(inspect_content(b"Public information",policy())["decision"],"allow")
+        self.assertEqual(inspect_content(b"Secret Roadmap",policy())["ruleId"],"secret")
+        self.assertEqual(inspect_content(b"\xff\x00",policy())["decision"],"block")
+    def test_office_extraction(self):
+        output=io.BytesIO()
+        with zipfile.ZipFile(output,"w") as archive:
+            archive.writestr("word/document.xml",'<document xmlns="urn:test"><t>Secret Roadmap</t></document>')
+        result=inspect_content(output.getvalue(),policy())
+        self.assertEqual(result["contentType"],"office");self.assertEqual(result["ruleId"],"secret")
+    def test_pdf_extraction(self):
+        from pypdf import PdfWriter
+        from pypdf.generic import DecodedStreamObject,NameObject,DictionaryObject
+        writer=PdfWriter();page=writer.add_blank_page(300,300)
+        font=DictionaryObject({NameObject('/Type'):NameObject('/Font'),NameObject('/Subtype'):NameObject('/Type1'),NameObject('/BaseFont'):NameObject('/Helvetica')})
+        page[NameObject('/Resources')]=DictionaryObject({NameObject('/Font'):DictionaryObject({NameObject('/F1'):writer._add_object(font)})})
+        content=DecodedStreamObject();content.set_data(b'BT /F1 12 Tf 10 100 Td (Secret Roadmap) Tj ET');page[NameObject('/Contents')]=writer._add_object(content)
+        output=io.BytesIO();writer.write(output)
+        result=inspect_content(output.getvalue(),policy());self.assertEqual(result['contentType'],'pdf');self.assertEqual(result['ruleId'],'secret')
+    def test_semantic_threshold_and_missing_model(self):
+        p=policy();p['dlp']['semantic']={'enabled':True,'model':'google/embeddinggemma-2','revision':'test-revision','categories':[{'id':'semantic-secret','threshold':.8,'prototype':[1.,0.]}]}
+        os.environ['EMBEDDING_MODEL_REVISION']='test-revision'
+        result=inspect_content(b'Confidential paraphrase',p,encoder=lambda *_:[1.,0.]);self.assertEqual(result['ruleId'],'semantic-secret')
+        result=inspect_content(b'Confidential paraphrase',p,encoder=lambda *_: (_ for _ in ()).throw(ValueError()));self.assertEqual(result['decision'],'block')
+    def test_cosine_validates(self):
+        self.assertAlmostEqual(cosine([1,0],[1,0]),1)
+        with self.assertRaises(ValueError):cosine([0,0],[0,0])
+    def test_video_has_no_false_safe_coverage(self):
+        video=b'\x00\x00\x00\x18ftypisom'+b'\x00'*20
+        self.assertEqual(content_kind(video),'video');result=inspect_content(video,policy());self.assertEqual(result['decision'],'block');self.assertFalse(result['coverage']['complete'])
+
+if __name__=='__main__':unittest.main()