airlock/tests/test_batch_file_state.py at main · coreyt/airlock · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
"""Unit tests for the file-scan state machine + async worker (to-do #2)."""

from __future__ import annotations

import json

from airlock.batch import scan, worker
from airlock.batch.store import (
    FILE_FAILED,
    FILE_READY,
    FILE_REJECTED,
    FILE_SCANNING,
    FILE_UPLOADED,
    BatchStore,
)


def _store(tmp_path):
    return BatchStore(str(tmp_path / "b.db"))


# ---------------------------------------------------------------------------
# Store: batch_files lifecycle
# ---------------------------------------------------------------------------
class TestFileStateStore:
    def test_record_then_get(self, tmp_path):
        s = _store(tmp_path)
        s.record_file_upload("file-1", byte_count=42)
        row = s.get_file("file-1")
        assert row["status"] == FILE_UPLOADED
        assert row["byte_count"] == 42

    def test_claim_is_won_once(self, tmp_path):
        s = _store(tmp_path)
        s.record_file_upload("file-1", byte_count=1)
        assert s.claim_file_scan("file-1") is True
        assert s.get_file("file-1")["status"] == FILE_SCANNING
        # Second claim while leased -> lost (scan runs once).
        assert s.claim_file_scan("file-1") is False

    def test_terminal_transitions(self, tmp_path):
        s = _store(tmp_path)
        s.record_file_upload("ready", byte_count=1)
        s.claim_file_scan("ready")
        s.set_file_ready("ready", row_count=7)
        assert s.get_file("ready")["status"] == FILE_READY
        assert s.get_file("ready")["row_count"] == 7

        s.record_file_upload("bad", byte_count=1)
        s.claim_file_scan("bad")
        s.set_file_rejected("bad", reason="blocked keyword: 'x'")
        assert s.get_file("bad")["status"] == FILE_REJECTED
        assert "blocked" in s.get_file("bad")["reason"]

    def test_claim_unknown_file_is_false(self, tmp_path):
        assert _store(tmp_path).claim_file_scan("nope") is False

    def test_scan_enabled_flag_is_persisted(self, tmp_path):
        s = _store(tmp_path)
        s.record_file_upload("scanned", byte_count=1)  # default True
        s.record_file_upload("raw", byte_count=1, status=FILE_READY, scan_enabled=False)
        assert s.get_file("scanned")["scan_enabled"] == 1
        assert s.get_file("raw")["scan_enabled"] == 0


# ---------------------------------------------------------------------------
# Worker: run_scan + await_file_ready (real thread pool, fake guards via profile)
# ---------------------------------------------------------------------------
def _write(tmp_path, name, text):
    p = tmp_path / name
    p.write_text(
        json.dumps(
            {
                "custom_id": "r1",
                "body": {"messages": [{"role": "user", "content": text}]},
            }
        )
        + "\n"
    )
    return p


class TestWorker:
    async def test_run_scan_marks_ready(self, tmp_path, monkeypatch):
        # No keywords / no PII -> clean passthrough, no Presidio.
        monkeypatch.delenv("AIRLOCK_BLOCKED_KEYWORDS", raising=False)
        s = _store(tmp_path)
        src = _write(tmp_path, "in.jsonl", "hello world")
        dst = tmp_path / "in.scrubbed.jsonl"
        s.record_file_upload("file-1", byte_count=src.stat().st_size)

        await worker.run_scan(
            s,
            "file-1",
            str(src),
            str(dst),
            {"scan_at_upload": True, "keyword_block": True, "pii_redact": False},
        )
        assert s.get_file("file-1")["status"] == FILE_READY
        assert dst.exists()

    async def test_run_scan_marks_rejected_on_keyword(self, tmp_path, monkeypatch):
        monkeypatch.setenv("AIRLOCK_BLOCKED_KEYWORDS", "classified")
        s = _store(tmp_path)
        src = _write(tmp_path, "in.jsonl", "this is classified material")
        dst = tmp_path / "in.scrubbed.jsonl"
        s.record_file_upload("file-1", byte_count=src.stat().st_size)

        await worker.run_scan(
            s,
            "file-1",
            str(src),
            str(dst),
            {"scan_at_upload": True, "keyword_block": True, "pii_redact": False},
        )
        assert s.get_file("file-1")["status"] == FILE_REJECTED
        assert not dst.exists()

    async def test_run_scan_failure_marks_failed(self, tmp_path, monkeypatch):
        s = _store(tmp_path)
        s.record_file_upload("file-1", byte_count=1)

        def boom(*a, **k):
            raise RuntimeError("executor blew up")

        monkeypatch.setattr(scan, "scan_file", boom)
        await worker.run_scan(s, "file-1", "x", "y", {})
        assert s.get_file("file-1")["status"] == FILE_FAILED
        assert "executor blew up" in s.get_file("file-1")["reason"]

    async def test_await_ready_returns_terminal_row(self, tmp_path):
        s = _store(tmp_path)
        s.record_file_upload("file-1", byte_count=1)
        s.claim_file_scan("file-1")
        s.set_file_ready("file-1", row_count=1)
        row = await worker.await_file_ready(s, "file-1", timeout=1.0)
        assert row["status"] == FILE_READY

    async def test_await_ready_times_out_while_scanning(self, tmp_path):
        s = _store(tmp_path)
        s.record_file_upload("file-1", byte_count=1)
        s.claim_file_scan("file-1")  # left SCANNING
        row = await worker.await_file_ready(s, "file-1", timeout=0.05, interval=0.01)
        assert row["status"] == FILE_SCANNING

    async def test_await_ready_unknown_file_is_none(self, tmp_path):
        assert (
            await worker.await_file_ready(_store(tmp_path), "x", timeout=0.05) is None
        )