88 lines
5.6 KiB
Python
88 lines
5.6 KiB
Python
import hashlib
|
|
import io
|
|
import json
|
|
import queue
|
|
import tempfile
|
|
from pathlib import Path
|
|
import unittest
|
|
from unittest.mock import patch
|
|
from semantic import NativeSemanticModel,validate_vectors,validate_readiness,verify_file,DIMENSIONS,PROFILE,RUNTIME_BUILD
|
|
from engine import aggregate_vectors,embedding_features,inspect_content
|
|
|
|
def vector(index=0):
|
|
result=[0.0]*DIMENSIONS;result[index]=1.0;return result
|
|
|
|
class SemanticProtocolTests(unittest.TestCase):
|
|
def model(self,windows):
|
|
model=NativeSemanticModel.__new__(NativeSemanticModel)
|
|
def request(payload,timeout=30):
|
|
if payload['op']=='windows':return {'windows':windows}
|
|
return {'vectors':[vector() for _ in payload['texts']]}
|
|
model._request=request
|
|
return model
|
|
def test_unicode_windows_cover_exact_source_and_tail(self):
|
|
text='Confidential café 中央 secret';encoded=text.encode()
|
|
windows=[{'text':text,'start_byte':0,'end_byte':len(encoded)}]
|
|
self.assertEqual(self.model(windows).text_vectors(text),[vector()])
|
|
windows[0]['end_byte']-=1
|
|
with self.assertRaises(ValueError):self.model(windows).text_vectors(text)
|
|
def test_missing_tail_gaps_and_substitution_are_refused(self):
|
|
for windows in [[{'text':'hello','start_byte':0,'end_byte':5}],
|
|
[{'text':'hello','start_byte':0,'end_byte':5},{'text':'world','start_byte':6,'end_byte':11}],
|
|
[{'text':'altered text','start_byte':0,'end_byte':11}]]:
|
|
with self.assertRaises(ValueError):self.model(windows).text_vectors('hello world')
|
|
def test_full_window_batch_is_not_silently_truncated(self):
|
|
text='a'*65
|
|
windows=[{'text':'a','start_byte':index,'end_byte':index+1} for index in range(65)]
|
|
with self.assertRaises(ValueError):self.model(windows).text_vectors(text)
|
|
def test_output_dimensions_finite_norm_and_aggregation(self):
|
|
self.assertEqual(validate_vectors([vector()]),[vector()])
|
|
for invalid in [[0.0]*768,[1.0]*767,[float('nan')]*768,[float('inf')]*768]:
|
|
with self.assertRaises(ValueError):validate_vectors([invalid])
|
|
aggregate=aggregate_vectors([vector(0),vector(1)])
|
|
self.assertAlmostEqual(sum(value*value for value in aggregate),1)
|
|
def test_only_exact_hardened_cpu_worker_build_is_accepted(self):
|
|
ready={'ready':True,'model':'google/embeddinggemma-2','dimensions':768,'runtime':RUNTIME_BUILD,'device':'cpu','vulkan_module_loaded':False}
|
|
validate_readiness(ready)
|
|
for changed in [{'runtime':'b11461'},{'runtime':'b11461-no-logits-v2'},{'device':'gpu'},{'vulkan_module_loaded':True},{'dimensions':128}]:
|
|
with self.assertRaises(ValueError):validate_readiness({**ready,**changed})
|
|
def test_artifact_hash_size_and_escape_checks(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root=Path(directory).resolve();file=root/'artifact';file.write_bytes(b'public model fixture')
|
|
digest=hashlib.sha256(file.read_bytes()).hexdigest()
|
|
self.assertEqual(verify_file(root,'artifact',file.stat().st_size,digest),file)
|
|
with self.assertRaises(ValueError):verify_file(root,'artifact',file.stat().st_size,'0'*64)
|
|
with self.assertRaises(ValueError):verify_file(root,'artifact',1,digest)
|
|
with self.assertRaises(ValueError):verify_file(root,'../escape',1,digest)
|
|
def test_training_features_return_only_vectors_and_metadata(self):
|
|
class FixtureModel:
|
|
def text_vectors(self,text):return [vector(),vector(1)]
|
|
def info(self):return {'model':'google/embeddinggemma-2','revision':'a'*40,'dimensions':768,'profile':PROFILE,'local':True}
|
|
with patch('engine.get_model',return_value=FixtureModel()):result=embedding_features(b'Internal secret not to retain in reference data')
|
|
self.assertTrue(result['coverage']['complete']);self.assertEqual(result['coverage']['textChunks'],2)
|
|
self.assertEqual(len(result['vectors']),2);self.assertNotIn('Internal secret',json.dumps(result))
|
|
def test_category_matching_compares_all_reference_prototypes(self):
|
|
policy={'dlp':{'maxBytes':1024,'rules':[],'semantic':{'enabled':True,'model':'google/embeddinggemma-2','revision':'a'*40,'categories':[{'id':'reference','prototype':vector(0),'prototypes':[vector(0),vector(1)],'threshold':.9}]}}}
|
|
with patch.dict('os.environ',{'EMBEDDING_MODEL_REVISION':'a'*40}):
|
|
result=inspect_content(b'Paraphrased protected concept',policy,encoder=lambda *_:vector(1))
|
|
self.assertEqual(result['ruleId'],'reference')
|
|
def test_worker_timeout_kills_and_sanitizes_errors(self):
|
|
model=NativeSemanticModel.__new__(NativeSemanticModel);model.responses=queue.Queue()
|
|
with patch.object(model,'close') as close:
|
|
with self.assertRaisesRegex(ValueError,'timed out'):model._receive(.001)
|
|
close.assert_called_once()
|
|
def test_worker_private_exception_text_is_never_returned(self):
|
|
import threading
|
|
class Process:
|
|
stdin=io.BytesIO()
|
|
def poll(self):return None
|
|
def kill(self):pass
|
|
model=NativeSemanticModel.__new__(NativeSemanticModel);model.process=Process();model.lock=threading.Lock()
|
|
model._start=lambda:None
|
|
model._receive=lambda timeout:{'ok':False,'device':'cpu','error':'PRIVATE DOC CONTENT'}
|
|
with patch.object(model,'close') as close:
|
|
with self.assertRaises(ValueError) as error:model._request({'op':'documents','texts':['public fixture']})
|
|
self.assertNotIn('PRIVATE',str(error.exception));close.assert_called_once()
|
|
|
|
if __name__=='__main__':unittest.main()
|