okami.mom/albedo

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

dawnimages: refuse a PNG or JPEG cut short where it is attached, so it never reaches history79b55b3

main
branchesmainall refs →
5.0 KiB125 linesraw
1"""Catch retained image leaks and premature release during nested delivery.
2
3The daemon cannot inspect archived captures or pause/refuse cells.finish, so
4these checks use the real kernel through the shared owner channel.
5"""
6
7import unittest
8
9from remote_kernel_test import OwnerChannel
10
11PNG = "iVBORw0KGgoAAAANSUhEUgAAAAIAAAADCAIAAAA2iEnWAAAAAElFTkSuQmCC"
12IMAGE = f"show_image(__import__('base64').b64decode({PNG!r}))\n"
13INSPECT = """
14import __main__ as kernel
15
16def capture_state(id):
17    capture = kernel.ARCHIVES[id]
18    return {'images': len(capture.images), 'text': output.read(id)}
19
20def finished_images():
21    return {id: len(done['images']) for id, done in kernel.FINISHED.items()}
22"""
23
24
25class CellImagesTest(unittest.TestCase):
26    def setUp(self):
27        self.channel = OwnerChannel([])
28        self.addCleanup(self.channel.close)
29        self.assertEqual(self.channel.recv()["type"], "ready")
30        self.execute("setup", INSPECT)
31
32    def execute(self, id, code):
33        self.channel.send({"type": "execute", "id": id, "code": code})
34        return self.channel.wait_for(
35            lambda frame: frame["type"] == "done" and frame["id"] == id
36        )
37
38    def state(self, id):
39        reply = self.channel.invoke("inspect-" + id, name="capture_state", args=[id])
40        self.assertTrue(reply["ok"])
41        return reply["value"]
42
43    def acknowledge(self, call, value=None, refused=False):
44        answer = (
45            {"ok": False, "code": "refused", "message": "finish refused"}
46            if refused
47            else {"ok": True, "value": value}
48        )
49        self.channel.send({"type": "reply", "id": call["id"], "value": answer})
50
51    def finish_call(self):
52        started = self.channel.wait_for(
53            lambda frame: frame.get("method") == "cells.started"
54        )
55        self.acknowledge(started)
56        return self.channel.wait_for(
57            lambda frame: frame.get("method") == "cells.finish"
58        )
59
60    def test_delivered_success_error_and_cancellation_release_images(self):
61        for status, ending in (
62            ("ok", ""),
63            ("error", "raise ValueError('execution failed')"),
64            ("interrupted", "raise __import__('asyncio').CancelledError()"),
65        ):
66            with self.subTest(status=status):
67                done = self.execute(status, IMAGE + "print('retained text')\n" + ending)
68                self.assertEqual(done["status"], status)
69                self.assertEqual(done["images"], [PNG])
70                state = self.state(status)
71                self.assertEqual(state["images"], 0)
72                self.assertIn("retained text", state["text"])
73
74    def test_nested_images_wait_for_acknowledgement_and_survive_refusal(self):
75        for refused in (False, True):
76            with self.subTest(refused=refused):
77                child = f"child-{refused}"
78                parent = "parent-" + child
79                self.channel.send(
80                    {
81                        "type": "execute",
82                        "id": parent,
83                        "code": "await cells.run('saved')",
84                    }
85                )
86                prepare = self.channel.wait_for(
87                    lambda frame: frame.get("method") == "cells.prepare"
88                )
89                self.acknowledge(
90                    prepare,
91                    {"id": child, "source": IMAGE + "raise ValueError('child failed')"},
92                )
93                finish = self.finish_call()
94                self.assertEqual(finish["args"]["outcome"]["status"], "error")
95                self.assertEqual(finish["args"]["outcome"]["images"], [PNG])
96                self.assertEqual(self.state(child)["images"], 1)
97                self.assertEqual(self.state(parent)["images"], 1)
98                self.acknowledge(finish, refused=refused)
99                done = self.channel.wait_for(
100                    lambda frame: frame["type"] == "done" and frame["id"] == parent
101                )
102                self.assertEqual(done["images"], [PNG])
103                self.assertEqual(self.state(child)["images"], int(refused))
104                self.assertEqual(self.state(parent)["images"], 0)
105
106    def test_replays_keep_images_within_a_budget_newest_first(self):
107        self.execute("budget", f"kernel.FINISHED_IMAGE_BYTES = {len(PNG) * 2}")
108        for id in ("first", "second", "third"):
109            self.assertEqual(self.execute(id, IMAGE)["images"], [PNG])
110        reply = self.channel.invoke("finished", name="finished_images")
111        self.assertEqual(
112            {id: reply["value"][id] for id in ("first", "second", "third")},
113            {"first": 0, "second": 1, "third": 1},
114        )
115        self.channel.send({"type": "execute", "id": "first", "code": "unused"})
116        replayed = self.channel.wait_for(
117            lambda frame: frame["type"] == "done" and frame["id"] == "first"
118        )
119        self.assertEqual(replayed["images"], [])
120        self.assertIn("1 image(s) not kept for a replayed result", replayed["output"])
121        self.assertIn("attached image/png", replayed["value"])
122
123
124if __name__ == "__main__":
125    unittest.main()