okami.mom/albedo

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

nekomimi.petfeat: multiple images, collapsed pastes, and an attachment strip in the prompt6e82811

main
branchesmainall refs →
5.1 KiB131 linesraw
1"""Saved Python state expires with its idle session, and only then.
2
3Expiry deletes files the daemon cannot recreate, so the bugs worth catching are
4the ones that delete too much (a live kernel's file) or nothing at all (a type
5mismatch once made the sweep a no-op), and only the real daemon exercises both.
6"""
7
8import os
9import time
10import unittest
11
12from harness import latest_user, wait_until
13from harness import Albedo, Provider, exclusive, python, text
14
15IDLE_SECONDS = 1
16EXPIRY_SECONDS = 3
17LOST = (
18    "<system-note>The python kernel got reset and all variables are lost</system-note>"
19)
20
21
22# exclusive: changes daemon idle/expiry sweep timing and tests startup cleanup
23@exclusive
24class StateExpiryTests(unittest.TestCase):
25    def setUp(self):
26        def script(request):
27            if request["input"][-1].get("role") == "user":
28                prompt = latest_user(request)
29                if prompt.startswith("remember"):
30                    return python("answer = 7")
31                if prompt.startswith("pin"):
32                    return python("job = run('sleep', '25')\nanswer = 8")
33            return text("ok")
34
35        self.provider = Provider(script)
36        self.addCleanup(self.provider.close)
37        self.app = Albedo(
38            self.provider,
39            protocol="responses",
40            prepare=lambda app: app.env.update(
41                ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS),
42                ALBEDO_STATE_EXPIRY_SECONDS=str(EXPIRY_SECONDS),
43            ),
44        )
45        self.app.__enter__()
46        self.addCleanup(self.app.__exit__, None, None, None)
47
48    def state(self, session):
49        return self.app.home / "kernels" / f"{session}.state"
50
51    def turn(self, session, prompt):
52        self.app.prompt(session, prompt).close()
53        self.app.idle(session)
54        return latest_user(self.provider.requests[-1]["request"])
55
56    def test_an_idle_session_loses_its_state_and_a_live_kernel_keeps_it(self):
57        idle, live = self.app.session(), self.app.session()
58        for session in (idle, live):
59            self.turn(session, "remember this")
60        wait_until(
61            lambda: self.state(idle).exists() and self.state(live).exists(),
62            15,
63            "idle kernels were never saved",
64        )
65        # a background job pins this session's kernel, so it is not idle
66        self.turn(live, "pin this")
67        wait_until(
68            lambda: not self.state(idle).exists(), 30, "an idle session kept its state"
69        )
70        self.assertTrue(self.state(live).exists(), "a live kernel lost its state")
71        self.assertEqual(self.turn(idle, "recall it"), "recall it" + LOST)
72
73    def test_old_orphan_state_goes_at_startup_and_a_young_one_stays(self):
74        kernels = self.app.home / "kernels"
75        kernels.mkdir(exist_ok=True)
76        old, young = kernels / "orphan-old.state", kernels / "orphan-young.state"
77        for path in (old, young):
78            path.write_bytes(b"x")
79        two_days_ago = time.time() - 2 * 86400
80        os.utime(old, (two_days_ago, two_days_ago))
81        self.app.restart()
82        self.assertFalse(old.exists(), "an old orphan survived startup")
83        self.assertTrue(young.exists(), "a young orphan was deleted")
84
85
86# exclusive: changes daemon retention/idle timing and restarts the daemon
87@exclusive
88class StateRetentionTests(unittest.TestCase):
89    def setUp(self):
90        self.provider = Provider(
91            lambda request: (
92                python("answer = 7")
93                if request["input"][-1].get("role") == "user"
94                else text("ok")
95            )
96        )
97        self.addCleanup(self.provider.close)
98        self.app = Albedo(
99            self.provider,
100            protocol="responses",
101            prepare=lambda app: app.env.update(
102                ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS),
103                ALBEDO_STATE_EXPIRY_SECONDS=str(10 * 86400),
104            ),
105        )
106        self.app.__enter__()
107        self.addCleanup(self.app.__exit__, None, None, None)
108
109    def test_saved_pastes_older_than_the_retention_are_deleted(self):
110        folder = self.app.home / "pastes" / "some-session"
111        folder.mkdir(parents=True)
112        old, young = folder / "old-1.md", folder / "young-1.md"
113        for path in (old, young):
114            path.write_text("pasted")
115        eleven_days_ago = time.time() - 11 * 86400
116        os.utime(old, (eleven_days_ago, eleven_days_ago))
117        self.app.restart()
118        wait_until(lambda: not old.exists(), 15, "an expired paste was kept")
119        self.assertTrue(young.exists(), "a paste inside its retention was deleted")
120
121    def test_a_session_inside_its_retention_keeps_an_old_file(self):
122        session = self.app.session()
123        self.app.prompt(session, "remember this").close()
124        self.app.idle(session)
125        state = self.app.home / "kernels" / f"{session}.state"
126        wait_until(state.exists, 15, "the idle kernel was never saved")
127        two_days_ago = time.time() - 2 * 86400
128        os.utime(state, (two_days_ago, two_days_ago))
129        self.app.restart()
130        time.sleep(3)  # the first periodic sweep runs within a second of startup
131        self.assertTrue(state.exists(), "a session inside its retention lost its state")