kernelbot/tests/test_kernelguard.py at 6a573fdba1830808e10c63dc712b86fee61d9dd3 · gpu-mode/kernelbot · 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
import pytest

from libkernelbot import kernelguard
from libkernelbot.consts import SubmissionMode
from libkernelbot.utils import KernelBotError


def test_should_precheck_submission_enabled_modes(monkeypatch):
    monkeypatch.setenv("KERNELGUARD_ENABLED", "1")

    assert kernelguard.should_precheck_submission(SubmissionMode.BENCHMARK)
    assert kernelguard.should_precheck_submission(SubmissionMode.PROFILE)
    assert kernelguard.should_precheck_submission(SubmissionMode.LEADERBOARD)
    assert kernelguard.should_precheck_submission(SubmissionMode.PRIVATE)
    assert not kernelguard.should_precheck_submission(SubmissionMode.TEST)


def test_should_precheck_submission_disabled(monkeypatch):
    monkeypatch.delenv("KERNELGUARD_ENABLED", raising=False)
    assert not kernelguard.should_precheck_submission(SubmissionMode.BENCHMARK)


def test_timeout_sec_defaults_to_five_minutes(monkeypatch):
    monkeypatch.delenv("KERNELGUARD_TIMEOUT_SEC", raising=False)
    assert kernelguard._timeout_sec() == 5 * 60


def test_timeout_sec_can_be_overridden(monkeypatch):
    monkeypatch.setenv("KERNELGUARD_TIMEOUT_SEC", "240")
    assert kernelguard._timeout_sec() == 240


def test_enforce_submission_precheck_rejects_filtered_code(monkeypatch):
    monkeypatch.setenv("KERNELGUARD_ENABLED", "1")
    monkeypatch.setattr(
        kernelguard,
        "analyze_submission",
        lambda code: {
            "classification": "hacked",
            "should_filter": True,
            "filter_reason": "high_critical",
            "matched_patterns": [{"pattern": "MODULE_MUTATION"}],
        },
    )

    with pytest.raises(
        kernelguard.KernelGuardRejected,
        match=(
            r"Submission rejected by KernelGuard pre-check "
            r"\(high_critical\)\. Matched rules: MODULE_MUTATION\."
        ),
    ):
        kernelguard.enforce_submission_precheck("print('hello')", "submission.py")


def test_enforce_submission_precheck_fail_closed(monkeypatch):
    monkeypatch.setenv("KERNELGUARD_ENABLED", "1")
    monkeypatch.delenv("KERNELGUARD_FAIL_OPEN", raising=False)

    def _raise(_code):
        raise RuntimeError("boom")

    monkeypatch.setattr(kernelguard, "analyze_submission", _raise)

    with pytest.raises(
        KernelBotError,
        match="KernelGuard pre-check is unavailable right now. Please try again later.",
    ):
        kernelguard.enforce_submission_precheck("print('hello')", "submission.py")


def test_enforce_submission_precheck_fail_open(monkeypatch):
    monkeypatch.setenv("KERNELGUARD_ENABLED", "1")
    monkeypatch.setenv("KERNELGUARD_FAIL_OPEN", "1")

    def _raise(_code):
        raise RuntimeError("boom")

    monkeypatch.setattr(kernelguard, "analyze_submission", _raise)

    assert kernelguard.enforce_submission_precheck("print('hello')", "submission.py") is None


def test_analyze_submission_uses_cli_path(monkeypatch):
    expected = {"classification": "valid", "should_filter": False, "matched_patterns": []}
    monkeypatch.setattr(kernelguard, "_analyze_with_cli", lambda code: expected)

    assert kernelguard.analyze_submission("print('hello')") is expected