@ajksunkang-aios/kgraph-linux-x64 0.1.2 → 0.1.3

package/lib/site-packages/trio/_core/_tests/test_instrumentation.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NoReturn
+
+import attrs
+import pytest
+
+from ... import _abc, _core
+from .tutil import check_sequence_matches
+
+if TYPE_CHECKING:
+    from collections.abc import Container, Iterable
+
+    from ...lowlevel import Task
+
+
+@attrs.define(eq=False, slots=False)
+class TaskRecorder(_abc.Instrument):
+    record: list[tuple[str, Task | None]] = attrs.Factory(list)
+
+    def before_run(self) -> None:
+        self.record.append(("before_run", None))
+
+    def task_scheduled(self, task: Task) -> None:
+        self.record.append(("schedule", task))
+
+    def before_task_step(self, task: Task) -> None:
+        assert task is _core.current_task()
+        self.record.append(("before", task))
+
+    def after_task_step(self, task: Task) -> None:
+        assert task is _core.current_task()
+        self.record.append(("after", task))
+
+    def after_run(self) -> None:
+        self.record.append(("after_run", None))
+
+    def filter_tasks(self, tasks: Container[Task]) -> Iterable[tuple[str, Task | None]]:
+        for item in self.record:
+            if item[0] in ("schedule", "before", "after") and item[1] in tasks:
+                yield item
+            if item[0] in ("before_run", "after_run"):
+                yield item
+
+
+def test_instruments(recwarn: object) -> None:
+    r1 = TaskRecorder()
+    r2 = TaskRecorder()
+    r3 = TaskRecorder()
+
+    task = None
+
+    # We use a child task for this, because the main task does some extra
+    # bookkeeping stuff that can leak into the instrument results, and we
+    # don't want to deal with it.
+    async def task_fn() -> None:
+        nonlocal task
+        task = _core.current_task()
+
+        for _ in range(4):
+            await _core.checkpoint()
+        # replace r2 with r3, to test that we can manipulate them as we go
+        _core.remove_instrument(r2)
+        with pytest.raises(KeyError):
+            _core.remove_instrument(r2)
+        # add is idempotent
+        _core.add_instrument(r3)
+        _core.add_instrument(r3)
+        for _ in range(1):
+            await _core.checkpoint()
+
+    async def main() -> None:
+        async with _core.open_nursery() as nursery:
+            nursery.start_soon(task_fn)
+
+    _core.run(main, instruments=[r1, r2])
+
+    # It sleeps 5 times, so it runs 6 times.  Note that checkpoint()
+    # reschedules the task immediately upon yielding, before the
+    # after_task_step event fires.
+    expected = (
+        [("before_run", None), ("schedule", task)]
+        + [("before", task), ("schedule", task), ("after", task)] * 5
+        + [("before", task), ("after", task), ("after_run", None)]
+    )
+    assert r1.record == r2.record + r3.record
+    assert task is not None
+    assert list(r1.filter_tasks([task])) == expected
+
+
+def test_instruments_interleave() -> None:
+    tasks = {}
+
+    async def two_step1() -> None:
+        tasks["t1"] = _core.current_task()
+        await _core.checkpoint()
+
+    async def two_step2() -> None:
+        tasks["t2"] = _core.current_task()
+        await _core.checkpoint()
+
+    async def main() -> None:
+        async with _core.open_nursery() as nursery:
+            nursery.start_soon(two_step1)
+            nursery.start_soon(two_step2)
+
+    r = TaskRecorder()
+    _core.run(main, instruments=[r])
+
+    expected = [
+        ("before_run", None),
+        ("schedule", tasks["t1"]),
+        ("schedule", tasks["t2"]),
+        {
+            ("before", tasks["t1"]),
+            ("schedule", tasks["t1"]),
+            ("after", tasks["t1"]),
+            ("before", tasks["t2"]),
+            ("schedule", tasks["t2"]),
+            ("after", tasks["t2"]),
+        },
+        {
+            ("before", tasks["t1"]),
+            ("after", tasks["t1"]),
+            ("before", tasks["t2"]),
+            ("after", tasks["t2"]),
+        },
+        ("after_run", None),
+    ]
+    print(list(r.filter_tasks(tasks.values())))
+    check_sequence_matches(list(r.filter_tasks(tasks.values())), expected)
+
+
+def test_null_instrument() -> None:
+    # undefined instrument methods are skipped
+    class NullInstrument(_abc.Instrument):
+        def something_unrelated(self) -> None:
+            pass  # pragma: no cover
+
+    async def main() -> None:
+        await _core.checkpoint()
+
+    _core.run(main, instruments=[NullInstrument()])
+
+
+def test_instrument_before_after_run() -> None:
+    record = []
+
+    class BeforeAfterRun(_abc.Instrument):
+        def before_run(self) -> None:
+            record.append("before_run")
+
+        def after_run(self) -> None:
+            record.append("after_run")
+
+    async def main() -> None:
+        pass
+
+    _core.run(main, instruments=[BeforeAfterRun()])
+    assert record == ["before_run", "after_run"]
+
+
+def test_instrument_task_spawn_exit() -> None:
+    record = []
+
+    class SpawnExitRecorder(_abc.Instrument):
+        def task_spawned(self, task: Task) -> None:
+            record.append(("spawned", task))
+
+        def task_exited(self, task: Task) -> None:
+            record.append(("exited", task))
+
+    async def main() -> Task:
+        return _core.current_task()
+
+    main_task = _core.run(main, instruments=[SpawnExitRecorder()])
+    assert ("spawned", main_task) in record
+    assert ("exited", main_task) in record
+
+
+# This test also tests having a crash before the initial task is even spawned,
+# which is very difficult to handle.
+def test_instruments_crash(caplog: pytest.LogCaptureFixture) -> None:
+    record = []
+
+    class BrokenInstrument(_abc.Instrument):
+        def task_scheduled(self, task: Task) -> NoReturn:
+            record.append("scheduled")
+            raise ValueError("oops")
+
+        def close(self) -> None:
+            # Shouldn't be called -- tests that the instrument disabling logic
+            # works right.
+            record.append("closed")  # pragma: no cover
+
+    async def main() -> Task:
+        record.append("main ran")
+        return _core.current_task()
+
+    r = TaskRecorder()
+    main_task = _core.run(main, instruments=[r, BrokenInstrument()])
+    assert record == ["scheduled", "main ran"]
+    # the TaskRecorder kept going throughout, even though the BrokenInstrument
+    # was disabled
+    assert ("after", main_task) in r.record
+    assert ("after_run", None) in r.record
+    # And we got a log message
+    assert caplog.records[0].exc_info is not None
+    exc_type, exc_value, _exc_traceback = caplog.records[0].exc_info
+    assert exc_type is ValueError
+    assert str(exc_value) == "oops"
+    assert "Instrument has been disabled" in caplog.records[0].message
+
+
+def test_instruments_monkeypatch() -> None:
+    class NullInstrument(_abc.Instrument):
+        pass
+
+    instrument = NullInstrument()
+
+    async def main() -> None:
+        record: list[Task] = []
+
+        # Changing the set of hooks implemented by an instrument after
+        # it's installed doesn't make them start being called right away
+        instrument.before_task_step = (  # type: ignore[method-assign]
+            record.append  # type: ignore[assignment] # append is pos-only
+        )
+
+        await _core.checkpoint()
+        await _core.checkpoint()
+        assert len(record) == 0
+
+        # But if we remove and re-add the instrument, the new hooks are
+        # picked up
+        _core.remove_instrument(instrument)
+        _core.add_instrument(instrument)
+        await _core.checkpoint()
+        await _core.checkpoint()
+        assert record.count(_core.current_task()) == 2
+
+        _core.remove_instrument(instrument)
+        await _core.checkpoint()
+        await _core.checkpoint()
+        assert record.count(_core.current_task()) == 2
+
+    _core.run(main, instruments=[instrument])
+
+
+def test_instrument_that_raises_on_getattr() -> None:
+    class EvilInstrument(_abc.Instrument):
+        def task_exited(self, task: Task) -> NoReturn:
+            raise AssertionError("this should never happen")  # pragma: no cover
+
+        @property
+        def after_run(self) -> NoReturn:
+            raise ValueError("oops")
+
+    async def main() -> None:
+        with pytest.raises(ValueError, match=r"^oops$"):
+            _core.add_instrument(EvilInstrument())
+
+        # Make sure the instrument is fully removed from the per-method lists
+        runner = _core.current_task()._runner
+        assert "after_run" not in runner.instruments
+        assert "task_exited" not in runner.instruments
+
+    _core.run(main)
+
+
+def test_instrument_call_trio_context() -> None:
+    called = set()
+
+    class Instrument(_abc.Instrument):
+        pass
+
+    hooks = {
+        # not run in task context
+        "after_io_wait": (True, False),
+        "before_io_wait": (True, False),
+        "before_run": (True, False),
+        "after_run": (True, False),
+        # run in task context
+        "before_task_step": (True, True),
+        "after_task_step": (True, True),
+        "task_exited": (True, True),
+        # depends
+        "task_scheduled": (True, None),
+        "task_spawned": (True, None),
+    }
+    for hook, val in hooks.items():
+
+        def h(
+            self: Instrument,
+            *args: object,
+            hook: str = hook,
+            val: tuple[bool, bool | None] = val,
+        ) -> None:
+            fail_str = f"failed in {hook}"
+
+            assert _core.in_trio_run() == val[0], fail_str
+            if val[1] is not None:
+                assert _core.in_trio_task() == val[1], fail_str
+            called.add(hook)
+
+        setattr(Instrument, hook, h)
+
+    async def main() -> None:
+        await _core.checkpoint()
+
+        async with _core.open_nursery() as nursery:
+            nursery.start_soon(_core.checkpoint)
+
+    _core.run(main, instruments=[Instrument()])
+    assert called == set(hooks)