safetynet/apps/inspector/test_engine.py

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("[Content_Types].xml",'<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/><Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>')
            archive.writestr("_rels/.rels",'<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/></Relationships>')
            archive.writestr("word/document.xml",'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>Secret Roadmap</w:t></w:r></w:p></w:body></w:document>')
        result=inspect_content(output.getvalue(),policy())
        self.assertEqual(result["contentType"],"office");self.assertEqual(result["ruleId"],"secret")
        self.assertTrue(result["coverage"]["complete"])
    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()