GDPRScanner/tests/test_stop_scan.py

"""
Stop must stop — regressions for "pressing Stop seems to restart the scan".

  - SSE replay: a stopped/finished scan's buffered events are never replayed to
    a reconnecting browser (that re-showed the Stop button and progress bar).
  - run_scan: Stop during collection ends the scan (checkpoint + scan_cancelled)
    instead of announcing scan_start / "Scanning…".
  - Scheduler: Stop during one stage prevents later stages — in particular the
    Google stage no longer clears the Google abort flag and scans everyone.
"""
from __future__ import annotations
from unittest.mock import MagicMock

import pytest

import sse


@pytest.fixture(autouse=True)
def clean_scan_state():
    """Restore SSE globals and abort flags so tests don't leak into each other."""
    from routes import state
    saved_buf = list(sse._sse_buffer)
    saved_id = sse._current_scan_id
    saved_finished = list(sse._finished_scan_ids)
    sse._sse_buffer.clear()
    sse._current_scan_id = ""
    sse._finished_scan_ids.clear()
    yield
    state._scan_abort.clear()
    state._google_scan_abort.clear()
    sse._sse_buffer.clear()
    sse._sse_buffer.extend(saved_buf)
    sse._current_scan_id = saved_id
    sse._finished_scan_ids.clear()
    sse._finished_scan_ids.extend(saved_finished)


def _events(msgs: list[str]) -> list[str]:
    return [m.split("\n", 1)[0].removeprefix("event: ") for m in msgs]


# ---------------------------------------------------------------------------
# SSE replay
# ---------------------------------------------------------------------------

class TestReplay:
    def test_running_scan_is_replayed(self):
        sse._current_scan_id = "scan_1"
        sse.broadcast("scan_phase", {"phase": "Collecting…"})
        sse.broadcast("scan_progress", {"pct": 10})
        assert _events(sse.replay_messages()) == ["scan_phase", "scan_progress"]

    def test_stopped_scan_is_not_replayed(self):
        sse._current_scan_id = "scan_1"
        sse.broadcast("scan_phase", {"phase": "Collecting…"})
        sse.broadcast("scan_start", {"total": 5})
        sse.broadcast("scan_progress", {"pct": 40})
        sse.broadcast("scan_cancelled", {"completed": 2})
        assert sse._current_scan_id == ""
        assert sse.replay_messages() == []

    def test_completed_scan_is_not_replayed(self):
        sse._current_scan_id = "scan_1"
        sse.broadcast("scan_progress", {"pct": 100})
        sse.broadcast("scan_done", {"total_scanned": 1, "flagged_count": 0})
        assert sse.replay_messages() == []

    def test_file_scan_done_ends_the_file_scan(self):
        # Previously only scan_done cleared the id, so every finished file scan
        # stayed "current" and was replayed on each page load.
        sse._current_scan_id = "filescan_1"
        sse.broadcast("scan_phase", {"phase": "Files — Share"})
        sse.broadcast("file_scan_done", {"total_scanned": 3, "flagged_count": 0})
        assert sse._current_scan_id == ""
        assert sse.replay_messages() == []

    def test_untagged_events_still_replayed_after_scan_ends(self):
        sse._current_scan_id = "scan_1"
        sse.broadcast("scan_progress", {"pct": 40})
        sse.broadcast("scan_cancelled", {"completed": 2})
        sse.broadcast("google_scan_done", {"total_scanned": 7})  # no scan running → untagged
        assert _events(sse.replay_messages()) == ["google_scan_done"]

    def test_new_scan_replays_only_its_own_events(self):
        sse._current_scan_id = "scan_1"
        sse.broadcast("scan_progress", {"pct": 40})
        sse.broadcast("scan_cancelled", {"completed": 2})
        sse._current_scan_id = "scan_2"
        sse.broadcast("scan_phase", {"phase": "Collecting…"})
        msgs = sse.replay_messages()
        assert _events(msgs) == ["scan_phase"]
        assert '"scan_id": "scan_2"' in msgs[0]


# ---------------------------------------------------------------------------
# run_scan: Stop during collection
# ---------------------------------------------------------------------------

class TestRunScanStopDuringCollection:
    @pytest.fixture()
    def engine(self, monkeypatch):
        import scan_engine
        from routes import state
        sent: list[tuple[str, dict]] = []
        checkpoints: list = []
        monkeypatch.setattr(scan_engine, "broadcast", lambda ev, data: sent.append((ev, data)))
        monkeypatch.setattr(scan_engine, "DB_OK", False)
        monkeypatch.setattr(scan_engine, "_load_checkpoint", lambda key, **kw: None)
        monkeypatch.setattr(scan_engine, "_save_checkpoint", lambda *a, **kw: checkpoints.append(a))
        monkeypatch.setattr(scan_engine, "_clear_checkpoint", lambda **kw: None)
        conn = MagicMock()
        conn.is_app_mode = True
        monkeypatch.setattr(state, "connector", conn)
        monkeypatch.setattr(state, "flagged_items", [])
        opts = {"sources": [], "user_ids": [{"id": "u1", "displayName": "User"}], "options": {}}
        return scan_engine, sent, checkpoints, opts

    def test_stop_before_processing_does_not_announce_scan_start(self, engine):
        scan_engine, sent, checkpoints, opts = engine
        from routes import state
        state._scan_abort.set()
        scan_engine.run_scan(opts)
        names = [e for e, _ in sent]
        assert "scan_cancelled" in names
        assert "scan_start" not in names
        assert not any(e == "scan_phase" and "Scanning" in d.get("phase", "") for e, d in sent)
        assert "scan_done" not in names
        assert len(checkpoints) == 1  # resumable

    def test_without_stop_scan_runs_to_done(self, engine):
        scan_engine, sent, _, opts = engine
        scan_engine.run_scan(opts)
        names = [e for e, _ in sent]
        assert "scan_start" in names and "scan_done" in names
        assert "scan_cancelled" not in names


# ---------------------------------------------------------------------------
# Scheduler: Stop ends the whole job
# ---------------------------------------------------------------------------

class TestSchedulerStop:
    JOB = {"id": "j1", "name": "Nightly", "profile_id": "", "enabled": True,
           "auto_email": True, "auto_retention": True, "retention_years": 5}
    OPTIONS = {"sources": ["email", "gmail"], "user_ids": [], "options": {},
               "file_sources": [{"id": "f1", "path": "/srv/share", "label": "Share"}],
               "google_sources": ["gmail"]}

    @pytest.fixture()
    def job(self, monkeypatch):
        import gdpr_scanner
        import scan_scheduler
        import routes.google_scan as google_scan
        from routes import state

        sent: list[tuple[str, dict]] = []
        monkeypatch.setattr(scan_scheduler, "load_jobs", lambda: [dict(self.JOB)])
        monkeypatch.setattr(gdpr_scanner, "_get_db", lambda *a, **kw: None)
        monkeypatch.setattr(gdpr_scanner, "broadcast", lambda ev, data: sent.append((ev, data)))
        conn = MagicMock()
        conn.is_authenticated.return_value = True
        monkeypatch.setattr(state, "connector", conn)
        monkeypatch.setattr(state, "google_connector", MagicMock())
        monkeypatch.setattr(state, "flagged_items", [{"id": "found"}])

        mocks = {
            "run_scan": MagicMock(),
            "run_file_scan": MagicMock(),
            "google": MagicMock(),
        }
        monkeypatch.setattr(gdpr_scanner, "run_scan", mocks["run_scan"])
        monkeypatch.setattr(gdpr_scanner, "run_file_scan", mocks["run_file_scan"])
        monkeypatch.setattr(google_scan, "_run_google_scan", mocks["google"])

        sched = scan_scheduler.ScanScheduler()
        monkeypatch.setattr(sched, "_build_options", lambda cfg: {
            k: (list(v) if isinstance(v, list) else dict(v) if isinstance(v, dict) else v)
            for k, v in self.OPTIONS.items()})
        mocks["email"] = MagicMock()
        mocks["retention"] = MagicMock()
        monkeypatch.setattr(sched, "_send_email_report", mocks["email"])
        monkeypatch.setattr(sched, "_run_retention", mocks["retention"])
        return sched, mocks, sent

    def test_stop_during_m365_stage_skips_file_and_google(self, job):
        from routes import state
        sched, mocks, sent = job

        def _user_presses_stop(opts):  # what POST /api/scan/stop does mid-scan
            state._scan_abort.set()
            state._google_scan_abort.set()
        mocks["run_scan"].side_effect = _user_presses_stop

        sched._execute_scan("j1")

        mocks["run_scan"].assert_called_once()
        mocks["run_file_scan"].assert_not_called()
        mocks["google"].assert_not_called()
        assert state._google_scan_abort.is_set(), "Google abort must not be cleared after Stop"
        run = sched._last_runs["j1"]
        assert run["status"] == "cancelled"
        mocks["email"].assert_not_called()
        mocks["retention"].assert_not_called()
        done = [d for e, d in sent if e == "scheduler_done"]
        assert done and done[0]["cancelled"] is True

    def test_uninterrupted_job_runs_every_stage(self, job):
        sched, mocks, sent = job
        sched._execute_scan("j1")
        mocks["run_scan"].assert_called_once()
        mocks["run_file_scan"].assert_called_once()
        mocks["google"].assert_called_once()
        assert sched._last_runs["j1"]["status"] == "completed"
        mocks["email"].assert_called_once()
        mocks["retention"].assert_called_once()
        done = [d for e, d in sent if e == "scheduler_done"]
        assert done and done[0]["cancelled"] is False