import os import pytest from ayn_antivirus.quarantine.vault import QuarantineVault def test_quarantine_and_restore(tmp_path): vault_dir = tmp_path / "vault" key_file = tmp_path / "keys" / "vault.key" vault = QuarantineVault(str(vault_dir), str(key_file)) test_file = tmp_path / "malware.txt" test_file.write_text("this is malicious content") threat_info = { "threat_name": "TestVirus", "threat_type": "virus", "severity": "high" } qid = vault.quarantine_file(str(test_file), threat_info) assert qid is not None assert not test_file.exists() assert vault.count() == 1 restore_path = tmp_path / "restored.txt" vault.restore_file(qid, str(restore_path)) assert restore_path.exists() assert restore_path.read_text() == "this is malicious content" def test_quarantine_list(tmp_path): vault_dir = tmp_path / "vault" key_file = tmp_path / "keys" / "vault.key" vault = QuarantineVault(str(vault_dir), str(key_file)) test_file = tmp_path / "test.txt" test_file.write_text("content") vault.quarantine_file(str(test_file), {"threat_name": "Test", "threat_type": "virus", "severity": "low"}) items = vault.list_quarantined() assert len(items) == 1 def test_quarantine_delete(tmp_path): vault_dir = tmp_path / "vault" key_file = tmp_path / "keys" / "vault.key" vault = QuarantineVault(str(vault_dir), str(key_file)) test_file = tmp_path / "test.txt" test_file.write_text("content") qid = vault.quarantine_file(str(test_file), {"threat_name": "Test", "threat_type": "virus", "severity": "low"}) assert vault.delete_file(qid) == True assert vault.count() == 0