okami.mom/albedo

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

Aly Raffaufrefactor(tests): share repeated E2E fixtures5bdc0ee

main
branchesmainall refs →
4.9 KiB119 linesraw
1"""A cell's output past the preview is spilled whole to a file, and a cell that
2pushes the kernel past its memory cap is interrupted while the namespace lives on."""
3
4import json
5import os
6import unittest
7
8from harness import Albedo, Provider, exclusive, latest_user, python, text
9
10
11LINES = 3000  # ~200 KiB, past the 64 KiB preview and the ascii shortcut
12
13# Grows by 8 MiB every 50 ms until the cap stops it; the chunks stay in the
14# namespace, so the kernel is still past the cap when the next cell starts.
15HOG = (
16    "import time\n"
17    "chunks = []\n"
18    "while True:\n"
19    "    chunks.append(bytearray(8 * 1024 * 1024))\n"
20    "    time.sleep(0.05)\n"
21)
22
23
24# exclusive: the memory cap is daemon environment
25@exclusive
26class CellLimitsTests(unittest.TestCase):
27    def setUp(self):
28        self.results = []
29
30        def script(request):
31            last = request["messages"][-1]
32            if last.get("role") == "tool":
33                self.results.append(json.loads(last["content"]))
34                return text("noted")
35            user = latest_user(request)
36            if "large spill" in user:
37                return python("print('p' * 2097152 + 'the end', end='')")
38            if "page spill" in user:
39                return python(
40                    f"id = {self.large_cell!r}\n"
41                    "assert output.read(id, offset=1048576-2, limit=10) == 'p' * 10\n"
42                    "assert output.read(id, offset=2097152, limit=65536) == 'the end'\n"
43                    "assert output.read(id, offset=2097159, limit=65536) == ''\n"
44                    "print('cell spill pages verified')"
45                )
46            if "read the spill" in user:
47                return python(
48                    f"print(files.read({self.spill!r}, start_line={LINES - 1}))"
49                )
50            if "spill" in user:
51                return python(
52                    f"for i in range({LINES}):\n"
53                    "    print(f'line {i:05d} ' + 'x' * 60)\n"
54                    "print('the last line')"
55                )
56            if "hog" in user:
57                return python(HOG)
58            if "free" in user:
59                return python("import gc\ndel chunks\ngc.collect()\nprint('freed')")
60            return text("done")
61
62        def prepare(app):
63            app.env["ALBEDO_KERNEL_MEMORY_BYTES"] = str(96 * 1024 * 1024)
64
65        self.provider = Provider(script)
66        self.addCleanup(self.provider.close)
67        self.app = Albedo(self.provider, prepare=prepare)
68        self.app.__enter__()
69        self.addCleanup(self.app.__exit__, None, None, None)
70        self.session = self.app.session()
71
72    def cell(self, prompt):
73        self.app.prompt(self.session, prompt).close()
74        self.app.idle(self.session)
75        return self.results[-1]
76
77    def test_output_past_the_preview_is_spilled_whole(self):
78        result = self.cell("spill")
79        self.assertEqual(result["status"], "ok")
80        self.assertTrue(result["truncated"])
81        self.assertIn("line 00000", result["output"])
82        self.assertNotIn("the last line", result["output"].split("[output was")[0])
83        note = result["output"].rsplit("[output was", 1)[1]
84        self.spill = note.split(" in ", 1)[1].split(":")[0]
85        self.assertEqual(
86            os.path.dirname(self.spill), str(self.app.home / "output"), self.spill
87        )
88        with open(self.spill, "rb") as spilled:
89            content = spilled.read()
90        self.assertTrue(content.startswith(b"line 00000 xxx"), content[:40])
91        self.assertTrue(content.endswith(b"the last line\n"), content[-40:])
92        self.assertTrue(
93            note.startswith(f" {len(content)} bytes; all of it is in"), note
94        )
95        read = self.cell("read the spill")
96        self.assertIn(f"line {LINES - 1:05d}", read["output"])
97        self.assertIn("the last line", read["output"])
98
99    def test_cell_output_pages_its_spill_past_the_memory_prefix(self):
100        spilled = self.cell("large spill")
101        self.assertEqual(spilled["status"], "ok", spilled)
102        self.large_cell = spilled["cell_id"]
103        result = self.cell("page spill")
104        self.assertEqual(result["status"], "ok", result)
105        self.assertIn("cell spill pages verified", result["output"])
106
107    def test_a_cell_past_the_memory_cap_is_interrupted_and_the_namespace_kept(self):
108        hog = self.cell("hog")
109        self.assertEqual(hog["status"], "interrupted", hog)
110        self.assertIn("kernel memory passed its cap", hog["output"])
111        self.assertIn("cap 96 MiB", hog["output"])
112        # Over the cap already, a cell that frees memory still runs.
113        freed = self.cell("free")
114        self.assertEqual(freed["status"], "ok", freed)
115        self.assertEqual(freed["output"].strip(), "freed")
116        # The memory really went: the next hog is stopped at the cap again.
117        again = self.cell("hog")
118        self.assertEqual(again["status"], "interrupted", again)
119        self.assertIn("kernel memory passed its cap", again["output"])