lazyBoy/tests/file-journal-capacity.test.py

"""Real SQLite admission pressure preserves receipts and reusable storage."""
import hashlib
from concurrent.futures import ThreadPoolExecutor
import json
from pathlib import Path
import sqlite3
import tempfile
import threading
import unittest


class FileJournalCapacity(unittest.TestCase):
    def setUp(self):
        source = Path(__file__).resolve().parents[1] / 'crates/supervisor/src/runner_files.py'
        self.module = {}
        exec(compile(source.read_text(), str(source), 'exec'), self.module)
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.db = sqlite3.connect(Path(self.temp.name) / 'journal.sqlite')
        self.addCleanup(self.db.close)
        self.module['file_journal_schema'](self.db)

    def test_unresolved_slots_include_unknown_and_release_without_deleting_receipts(self):
        self.db.executemany('INSERT INTO file_operations VALUES (?,?,?,?,?,?,?)',
                            [(str(i), '{}', '', 'unknown' if i % 2 else 'accepted', None, '{}', 0)
                             for i in range(1024)])
        self.db.commit()
        with self.assertRaisesRegex(ValueError, 'FILE_JOURNAL_QUOTA'):
            with self.db:
                self.db.execute('BEGIN IMMEDIATE')
                self.module['file_journal_admit'](self.db, {'bot_id': 'capacity'})
        with self.db:
            self.db.execute("UPDATE file_operations SET status='succeeded', result='{}' WHERE id='0'")
        with self.db:
            self.db.execute('BEGIN IMMEDIATE')
            self.module['file_journal_admit'](self.db, {'bot_id': 'capacity'})
        self.assertEqual(self.db.execute('SELECT count(*) FROM file_operations').fetchone()[0], 1024)

    def test_storage_pressure_blocks_only_new_intents_and_reuses_freed_pages(self):
        identity = {'computer_id': 'fixture', 'bot_id': 'capacity', 'generation': 1}
        body = {'operation_id': 'known', 'action': 'write', 'path': 'file',
                'arguments': {'contentBase64': '', 'expectedHash': ''}}
        encoded = self.module['file_identity'](identity)
        key = hashlib.sha256((encoded + ':known').encode()).hexdigest()
        digest = hashlib.sha256(json.dumps(body, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
        receipt = {'status': 'failed', 'code': 73, 'stdout': '', 'stderr': 'CONFLICT: preserved receipt'}
        self.db.execute('INSERT INTO file_operations VALUES (?,?,?,?,?,?,?)',
                        (key, encoded, digest, 'failed', json.dumps(receipt), '{}', 0))
        # Actual allocated database pages, rather than a mocked quota result.
        self.db.execute('CREATE TABLE pressure (data BLOB)')
        self.db.execute('INSERT INTO pressure VALUES (zeroblob(?))', (128 * 1024 * 1024,))
        self.db.commit()
        original_pages = self.db.execute('PRAGMA page_count').fetchone()[0]
        self.module['file_process_ref'] = lambda pid: None
        def call(payload, action='file_mutate'):
            return self.module['file_journal_handle'](
                self.db, {'identity': identity, 'action': action, 'input': json.dumps(payload)},
                'raise AssertionError("quota must reject before helper launch")', lambda *_: False)
        with self.assertRaisesRegex(ValueError, 'FILE_JOURNAL_QUOTA'):
            call({**body, 'operation_id': 'new'})
        self.assertEqual(self.db.execute('SELECT count(*) FROM file_operations').fetchone()[0], 1)
        self.assertEqual(self.db.execute('SELECT count(*) FROM file_outbox').fetchone()[0], 0)
        self.assertEqual(call(body), receipt)
        self.assertEqual(call({'operation_id': 'known'}, 'file_status'), receipt)
        with self.db:
            self.db.execute('DELETE FROM pressure')
        self.assertEqual(self.db.execute('PRAGMA page_count').fetchone()[0], original_pages)
        self.assertGreater(self.db.execute('PRAGMA freelist_count').fetchone()[0], 0)
        with self.db:
            self.db.execute('BEGIN IMMEDIATE')
            self.module['file_journal_admit'](self.db, {'bot_id': 'capacity'})
        self.assertEqual(call(body), receipt)


    def test_concurrent_admissions_cannot_take_the_same_last_slot(self):
        for count, owner, error in [(1023, '{}', 'FILE_JOURNAL_QUOTA'),
                                    (255, '{"bot_id":"capacity"}', 'FILE_AGENT_QUOTA')]:
            with self.subTest(quota=error):
                self.db.execute('DELETE FROM file_operations')
                self.db.executemany('INSERT INTO file_operations VALUES (?,?,?,?,?,?,?)',
                                    [(str(i), owner, '', 'accepted', None, '{}', 0) for i in range(count)])
                self.db.commit()
                ready = threading.Barrier(2)
                def admit(operation):
                    with sqlite3.connect(Path(self.temp.name) / 'journal.sqlite', timeout=5) as db:
                        ready.wait(timeout=5)
                        try:
                            db.execute('BEGIN IMMEDIATE')
                            self.module['file_journal_admit'](db, {'bot_id': 'capacity'})
                            db.execute('INSERT INTO file_operations VALUES (?,?,?,?,?,?,?)',
                                       (operation, owner, '', 'accepted', None, '{}', 0))
                            db.commit()
                            return 'accepted'
                        except ValueError as failure:
                            db.rollback()
                            return str(failure)
                with ThreadPoolExecutor(max_workers=2) as pool:
                    results = list(pool.map(admit, ['first', 'second']))
                self.assertCountEqual(results, ['accepted', error])
                self.assertEqual(self.db.execute('SELECT count(*) FROM file_operations').fetchone()[0], count + 1)


if __name__ == '__main__':
    unittest.main()