Viewing files at commit 3354584a5d91Commit detailsView latest

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