writeragent/tests/scripting/test_venv_probe_progress.py at master · KeithCu/writeragent · 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
"""Incremental venv self-check progress reporting."""

from unittest.mock import MagicMock, patch

from plugin.scripting.config_limits import SELF_CHECK_IMPORT_PROBE_TIMEOUT_SEC
from plugin.scripting.venv_diagnostics import run_venv_self_check_with_progress


def test_run_venv_self_check_with_progress_emits_grouped_present_missing() -> None:
    displays: list[str] = []
    statuses: list[str] = []
    execute_timeouts: list[int] = []

    def fake_execute(_self, script, timeout_sec=10):
        execute_timeouts.append(timeout_sec)
        if "platform" in script:
            return {"status": "ok", "result": {"v": "3.12.0", "arch": "x86_64"}}
        if "numpy" in script:
            return {"status": "ok", "result": "present"}
        return {"status": "ok", "result": None}

    mock_mgr = MagicMock()
    mock_mgr.execute.side_effect = lambda script, timeout_sec=10: fake_execute(None, script, timeout_sec)

    with (
        patch("plugin.scripting.venv_worker.PythonWorkerManager.get", return_value=mock_mgr),
        patch(
            "plugin.scripting.venv_diagnostics._probe_nlp_packages",
            return_value=({"spacy": "present", "textdescriptives": None, "transformers": None}, None),
        ),
        patch("plugin.scripting.venv_diagnostics.probe_vision_packages", return_value=({"docling": "present"}, None)),
        patch(
            "plugin.scripting.venv_diagnostics._probe_vector_search_packages",
            return_value=({"envwrap": "present", "sqlite_vec": "present"}, None),
        ),
    ):
        ok, msg = run_venv_self_check_with_progress(
            "/fake/python",
            displays.append,
            timeout=60.0,
            on_status=statuses.append,
            extra_lines_after_header=("Cython Accelerator: Active (Optimized)",),
        )

    assert ok is True
    assert "responds OK" in msg
    assert "Scientific Libraries: numpy" in msg
    assert "Missing:" in msg
    assert "Text / NLP Libraries" in msg
    assert "spacy" in msg
    assert "Cython Accelerator: Active (Optimized)" in msg
    assert any("Scientific Libraries: numpy" in text for text in displays)
    assert any("numpy" in text for text in displays)
    assert any("Vision" in text for text in displays)
    assert any("Vector Search" in text for text in displays)
    assert any("docling" in text for text in displays)
    assert any("envwrap" in text for text in displays)
    assert not any("... OK" in text for text in displays)
    assert any("numpy" in status for status in statuses)
    assert execute_timeouts
    assert all(timeout == SELF_CHECK_IMPORT_PROBE_TIMEOUT_SEC for timeout in execute_timeouts)


def test_run_venv_self_check_with_progress_continues_after_sandbox_probe_error() -> None:
    displays: list[str] = []
    call_count = {"n": 0}

    def fake_execute(_self, script, timeout_sec=10):
        if "platform" in script:
            return {"status": "ok", "result": {"v": "3.12.0", "arch": "x86_64"}}
        call_count["n"] += 1
        if call_count["n"] == 1:
            return {"status": "error", "message": "Python worker failed: timed out after 30 seconds"}
        return {"status": "ok", "result": None}

    mock_mgr = MagicMock()
    mock_mgr.execute.side_effect = lambda script, timeout_sec=10: fake_execute(None, script, timeout_sec)

    with (
        patch("plugin.scripting.venv_worker.PythonWorkerManager.get", return_value=mock_mgr),
        patch("plugin.scripting.venv_diagnostics._probe_nlp_packages", return_value=({}, None)),
        patch("plugin.scripting.venv_diagnostics.probe_vision_packages", return_value=({"docling": "present"}, None)),
        patch("plugin.scripting.venv_diagnostics._probe_vector_search_packages", return_value=({}, None)),
    ):
        ok, msg = run_venv_self_check_with_progress("/fake/python", displays.append, timeout=60.0)

    assert ok is True
    assert "Vision Libraries" in msg
    assert "Warning:" in msg


def test_run_venv_self_check_with_progress_skips_audio_and_vector_search() -> None:
    displays: list[str] = []

    def fake_execute(_self, script, timeout_sec=10):
        if "platform" in script:
            return {"status": "ok", "result": {"v": "3.12.0", "arch": "x86_64"}}
        return {"status": "ok", "result": None}

    mock_mgr = MagicMock()
    mock_mgr.execute.side_effect = lambda script, timeout_sec=10: fake_execute(None, script, timeout_sec)

    with (
        patch("plugin.scripting.venv_worker.PythonWorkerManager.get", return_value=mock_mgr),
        patch("plugin.scripting.venv_diagnostics._probe_nlp_packages", return_value=({}, None)),
        patch("plugin.scripting.venv_diagnostics.probe_vision_packages", return_value=({"docling": "present"}, None)),
        patch("plugin.scripting.venv_diagnostics._probe_audio_packages") as mock_audio,
        patch("plugin.scripting.venv_diagnostics._probe_vector_search_packages") as mock_vector,
    ):
        ok, msg = run_venv_self_check_with_progress(
            "/fake/python",
            displays.append,
            timeout=60.0,
            include_vector_search=False,
            include_audio=False,
        )

    assert ok is True
    mock_audio.assert_not_called()
    mock_vector.assert_not_called()
    assert "Vector Search" not in msg
    assert "Audio Recording" not in msg
    assert "Vision Libraries" in msg