okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffauffix(kernel): preserve statement identity during definition capture1c00b60

main
branchesmainall refs →
3.9 KiB110 linesraw
1"""A restarted kernel gets back its helpers, not only its variables.
2
3Saved state is written when an idle kernel is released and read when the next
4turn opens a new one, so these scenarios run through the real daemon: the
5restore path lives in the kernel and the notice in the daemon.
6"""
7
8import os
9import pickle
10import unittest
11
12from harness import latest_user, wait_until
13from harness import Albedo, Provider, exclusive, python, text
14
15IDLE_SECONDS = 1
16BIG_DEFINITION = "def big():\n    return '" + "x" * 40000 + "'\n"
17
18DEFINE = (
19    "import json\n"
20    "import os.path\n"
21    "import xml.dom\n"
22    "def helper(x):\n    return x * 2\n"
23    "class Box:\n    def __init__(self, v):\n        self.v = v\n"
24    "answer = 7\n"
25    "def rebound():\n    return 1\n"
26    "if answer:\n    rebound = 5\n" + BIG_DEFINITION + "blob = bytes(9 * 1024 * 1024)\n"
27)
28FAIL = (
29    "def survivor():\n    return 'alive'\n"
30    "pass\n"
31    "raise RuntimeError('boom')\n"
32    "def never_ran():\n    return 'unreachable'\n"
33)
34RECALL = (
35    "print('RESULT', helper(21), Box(3).v, json.dumps([1]), answer, rebound,\n"
36    "      survivor(), os.path.basename('a/b'), xml.dom.__name__,\n"
37    "      'big' in globals(), 'blob' in globals(), 'never_ran' in globals())"
38)
39LEGACY = "print('LEGACY', legacy)"
40
41
42# exclusive: requires the daemon-wide idle-kernel timing knob
43@exclusive
44class KernelStateTests(unittest.TestCase):
45    def setUp(self):
46        cells = {"define": DEFINE, "fail": FAIL, "recall": RECALL, "legacy": LEGACY}
47
48        def script(request):
49            if request["input"][-1].get("role") == "user":
50                prompt = latest_user(request).split("<system-note>")[0]
51                if prompt in cells:
52                    return python(cells[prompt])
53            return text("ok")
54
55        self.provider = Provider(script)
56        self.addCleanup(self.provider.close)
57        self.app = Albedo(
58            self.provider,
59            protocol="responses",
60            prepare=lambda app: app.env.update(ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS)),
61        )
62        self.app.__enter__()
63        self.addCleanup(self.app.__exit__, None, None, None)
64
65    def state(self, session):
66        return self.app.home / "kernels" / f"{session}.state"
67
68    def turn(self, session, prompt):
69        self.app.prompt(session, prompt).close()
70        self.app.idle(session)
71
72    def saved_after(self, session, since):
73        """The idle kernel was released again: its state file is newer than `since`."""
74        wait_until(
75            lambda: (
76                self.state(session).exists()
77                and self.state(session).stat().st_mtime_ns > since
78            ),
79            15,
80            "the idle kernel was not saved",
81        )
82
83    def outputs(self):
84        return "".join(
85            item.get("output", "")
86            for item in self.provider.requests[-1]["request"]["input"]
87        )
88
89    def test_helpers_classes_imports_and_failed_cells_survive_a_restart(self):
90        session = self.app.session()
91        self.turn(session, "define")
92        self.saved_after(session, 0)
93        saved = self.state(session).stat().st_mtime_ns
94        self.turn(session, "fail")
95        self.saved_after(session, saved)
96        self.turn(session, "recall")
97        output = self.outputs()
98        self.assertIn("RESULT 42 3 [1] 7 5 alive b xml.dom False False False", output)
99        notice = latest_user(self.provider.requests[-2]["request"])
100        self.assertIn("re-run from source", notice)
101        self.assertIn("helper", notice)
102
103    def test_a_state_file_in_the_old_format_still_restores_its_variables(self):
104        session = self.app.session()
105        self.turn(session, "define")
106        self.saved_after(session, 0)
107        legacy = {"cwd": os.getcwd(), "names": {"legacy": pickle.dumps(5)}}
108        self.state(session).write_bytes(pickle.dumps(legacy))
109        self.turn(session, "legacy")
110        self.assertIn("LEGACY 5", self.outputs())