File size: 2,593 Bytes
201dc77
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""U02-1 / U11-2: the triton_fa2 gate must not be written when no GPU test ran."""
import importlib.util
import os
import subprocess

import pytest

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _load():
    spec = importlib.util.spec_from_file_location(
        "gate_flash_attention_under_test", os.path.join(REPO, "scripts", "gate_flash_attention.py")
    )
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def _junit(passed=0, skipped=0, failed=0):
    cases = (
        "".join(f'<testcase classname="t" name="p{i}"/>' for i in range(passed))
        + "".join(f'<testcase classname="t" name="s{i}"><skipped message="x"/></testcase>' for i in range(skipped))
        + "".join(f'<testcase classname="t" name="f{i}"><failure message="x"/></testcase>' for i in range(failed))
    )
    return f'<?xml version="1.0"?><testsuites><testsuite name="pytest">{cases}</testsuite></testsuites>'


@pytest.fixture
def harness(monkeypatch):
    mod = _load()
    written = []
    monkeypatch.setattr(mod.gate_mod, "write_equivalence_gate", lambda **kw: written.append(kw) or "/x")
    monkeypatch.setattr(mod.gate_mod, "gate_status", lambda *a, **k: (True, "ok"))
    monkeypatch.setattr(mod.gate_mod, "kernel_source_hash", lambda: "h")

    def run_with(xml, cvd="0"):
        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cvd)

        def fake_run(cmd, **kw):
            arg = next(a for a in cmd if a.startswith("--junitxml="))
            with open(arg.split("=", 1)[1], "w") as fh:
                fh.write(xml)
            return subprocess.CompletedProcess(cmd, 0)

        monkeypatch.setattr(mod.subprocess, "run", fake_run)
        return mod.main()

    return run_with, written


def test_all_skipped_suite_does_not_write_gate(harness):
    run, written = harness
    assert run(_junit(skipped=35)) != 0
    assert written == []


def test_any_skip_does_not_write_gate(harness):
    run, written = harness
    assert run(_junit(passed=34, skipped=1)) != 0
    assert written == []


def test_fully_executed_suite_writes_gate(harness):
    run, written = harness
    assert run(_junit(passed=35)) == 0
    assert len(written) == 1


@pytest.mark.parametrize("cvd", ["0,1", "0,3", ",", "0,"])
def test_multiple_or_malformed_device_list_refused(harness, cvd):
    run, written = harness
    assert run(_junit(passed=35), cvd=cvd) == 2
    assert written == []


def test_unparseable_junit_report_does_not_write_gate(harness):
    run, written = harness
    assert run("not xml at all") != 0
    assert written == []