Viewing files at commit b4b51d4daa88Commit detailsView latest

Browse files
tests/test_runlock.py 153 lines · 5.4 KB · python
1
from __future__ import annotations
2
3
import contextlib
4
import os
5
import selectors
6
import subprocess
7
import sys
8
import tempfile
9
import textwrap
10
import time
11
import unittest
12
from pathlib import Path
13
from unittest import mock
14
15
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
16
17
from oakbench.runlock import measurement_lock
18
19
try:
20
    import fcntl  # noqa: F401
21
except ImportError:  # pragma: no cover - Windows
22
    HAS_FCNTL = False
23
else:
24
    HAS_FCNTL = True
25
26
27
SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts"
28
29
30
@contextlib.contextmanager
31
def isolated_lock_env(lock_path: Path, lock_value: str | None = None):
32
    updates = {"OAK_BENCH_LOCK_PATH": str(lock_path)}
33
    if lock_value is not None:
34
        updates["OAK_BENCH_LOCK"] = lock_value
35
    with mock.patch.dict(os.environ, updates, clear=False):
36
        if lock_value is None:
37
            os.environ.pop("OAK_BENCH_LOCK", None)
38
        yield
39
40
41
def subprocess_env(lock_path: Path) -> dict[str, str]:
42
    env = os.environ.copy()
43
    env["OAK_BENCH_LOCK_PATH"] = str(lock_path)
44
    env.pop("OAK_BENCH_LOCK", None)
45
    env["PYTHONDONTWRITEBYTECODE"] = "1"
46
    pythonpath = env.get("PYTHONPATH")
47
    env["PYTHONPATH"] = str(SCRIPTS_DIR) if not pythonpath else str(SCRIPTS_DIR) + os.pathsep + pythonpath
48
    return env
49
50
51
def readline_with_timeout(stream, timeout_s: float) -> str | None:
52
    selector = selectors.DefaultSelector()
53
    try:
54
        selector.register(stream, selectors.EVENT_READ)
55
        if not selector.select(timeout_s):
56
            return None
57
        return stream.readline()
58
    finally:
59
        selector.close()
60
61
62
class MeasurementLockTest(unittest.TestCase):
63
    @unittest.skipUnless(HAS_FCNTL, "fcntl not available")
64
    def test_acquiring_while_another_process_holds_lock_blocks_until_release(self) -> None:
65
        with tempfile.TemporaryDirectory(prefix="oakbench-runlock-test-") as tmp:
66
            tmp_path = Path(tmp)
67
            lock_path = tmp_path / "measurement.lock"
68
            holder_script = tmp_path / "holder.py"
69
            holder_script.write_text(
70
                textwrap.dedent(
71
                    f"""
72
                    import sys
73
                    import time
74
75
                    sys.path.insert(0, {str(SCRIPTS_DIR)!r})
76
77
                    from oakbench.runlock import measurement_lock
78
79
                    with measurement_lock("holder-process"):
80
                        print("ready", flush=True)
81
                        time.sleep(0.75)
82
                    """
83
                )
84
            )
85
            proc = subprocess.Popen(
86
                [sys.executable, str(holder_script)],
87
                stdout=subprocess.PIPE,
88
                stderr=subprocess.PIPE,
89
                text=True,
90
                env=subprocess_env(lock_path),
91
            )
92
            try:
93
                ready = readline_with_timeout(proc.stdout, 2.0) if proc.stdout is not None else None
94
                if ready != "ready\n":
95
                    proc.kill()
96
                    _, stderr = proc.communicate(timeout=2.0)
97
                    self.fail(f"holder process did not acquire lock: ready={ready!r} stderr={stderr}")
98
99
                start = time.perf_counter()
100
                with isolated_lock_env(lock_path):
101
                    with measurement_lock("waiter") as lock_info:
102
                        pass
103
                elapsed_ms = (time.perf_counter() - start) * 1000
104
105
                self.assertTrue(lock_info.enabled)
106
                self.assertIsNotNone(lock_info.wait_ms)
107
                self.assertGreaterEqual(lock_info.wait_ms, 400.0)
108
                self.assertGreaterEqual(elapsed_ms, 400.0)
109
110
                _, stderr = proc.communicate(timeout=2.0)
111
                self.assertEqual(proc.returncode, 0, stderr)
112
            finally:
113
                if proc.poll() is None:
114
                    proc.kill()
115
                    proc.communicate(timeout=2.0)
116
117
    def test_oak_bench_lock_off_bypasses(self) -> None:
118
        with tempfile.TemporaryDirectory(prefix="oakbench-runlock-test-") as tmp:
119
            lock_path = Path(tmp) / "measurement.lock"
120
            with isolated_lock_env(lock_path, "off"):
121
                with measurement_lock("disabled") as lock_info:
122
                    pass
123
124
            self.assertFalse(lock_info.enabled)
125
            self.assertIsNone(lock_info.wait_ms)
126
            self.assertFalse(lock_path.exists())
127
128
    @unittest.skipUnless(HAS_FCNTL, "fcntl not available")
129
    def test_wait_ms_is_near_zero_when_uncontended(self) -> None:
130
        with tempfile.TemporaryDirectory(prefix="oakbench-runlock-test-") as tmp:
131
            lock_path = Path(tmp) / "measurement.lock"
132
            with isolated_lock_env(lock_path):
133
                with measurement_lock("uncontended") as lock_info:
134
                    pass
135
136
            self.assertTrue(lock_info.enabled)
137
            self.assertIsNotNone(lock_info.wait_ms)
138
            self.assertLess(lock_info.wait_ms, 100.0)
139
140
    @unittest.skipUnless(HAS_FCNTL, "fcntl not available")
141
    def test_lock_file_content_contains_holder_label(self) -> None:
142
        with tempfile.TemporaryDirectory(prefix="oakbench-runlock-test-") as tmp:
143
            lock_path = Path(tmp) / "measurement.lock"
144
            with isolated_lock_env(lock_path):
145
                with measurement_lock("content-label"):
146
                    content = lock_path.read_text()
147
148
            self.assertIn("content-label", content)
149
            self.assertIn(f"pid={os.getpid()}", content)
150
151
152
if __name__ == "__main__":
153
    unittest.main()