broker/tests/test_store.py

import base64
import contextlib
import io
import os
import sqlite3
import sys
import tempfile
import unittest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)))

import coreci_broker

TABLE_COLUMNS = {
    "broker_config": ["key", "value"],
    "app_registrations": ["uuid", "name", "status", "created_at", "revoked_at"],
    "user_allowlist": ["user_id", "email", "status", "created_at", "revoked_at"],
    "agent_registrations": [
        "agent_id",
        "app_registration_uuid",
        "user_id",
        "tags",
        "token_hash",
        "token_expires_at",
        "first_registered_at",
        "last_seen_at",
    ],
    "admin_sessions": ["session_id", "created_at", "expires_at"],
}


class StoreTests(unittest.TestCase):
    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.env_path = os.path.join(tmp.name, ".env")
        self.db_path = os.path.join(tmp.name, "broker.db")

    def _boot(self, extra_env=None):
        env = {"BROKER_ENV_PATH": self.env_path, "BROKER_DB_PATH": self.db_path}
        if extra_env:
            env.update(extra_env)
        buffer = io.StringIO()
        with contextlib.redirect_stdout(buffer):
            broker = coreci_broker.boot(env)
        self.addCleanup(broker.store.close)
        return broker, buffer.getvalue()

    def _connect(self):
        conn = sqlite3.connect(self.db_path)
        self.addCleanup(conn.close)
        return conn

    def test_boot_creates_db_with_mode_0600(self):
        broker, stdout = self._boot()
        self.assertTrue(os.path.isfile(self.db_path))
        self.assertEqual(0o600, os.stat(self.db_path).st_mode & 0o777)
        self.assertEqual(1, stdout.count(broker.config["BROKER_ADMIN_PASS"]))

    def test_db_mode_0600_independent_of_umask(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        env = {
            "BROKER_ENV_PATH": os.path.join(tmp.name, ".env"),
            "BROKER_DB_PATH": os.path.join(tmp.name, "broker.db"),
        }
        old_umask = os.umask(0o777)
        try:
            with contextlib.redirect_stdout(io.StringIO()):
                broker = coreci_broker.boot(env)
            self.addCleanup(broker.store.close)
            self.assertEqual(0o600, os.stat(env["BROKER_DB_PATH"]).st_mode & 0o777)
            self.assertEqual(0o600, os.stat(env["BROKER_ENV_PATH"]).st_mode & 0o777)
        finally:
            os.umask(old_umask)

    def test_all_five_tables_present_with_columns(self):
        self._boot()
        conn = self._connect()
        tables = {
            row[0]
            for row in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")
        }
        self.assertTrue(set(TABLE_COLUMNS).issubset(tables))
        for table, expected in TABLE_COLUMNS.items():
            actual = [row[1] for row in conn.execute("PRAGMA table_info(%s)" % table)]
            self.assertEqual(expected, actual)

    def test_index_on_agent_registrations_app_registration_uuid(self):
        self._boot()
        conn = self._connect()
        names = [
            row[0]
            for row in conn.execute(
                "SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'agent_registrations'"
            )
        ]
        matching = []
        for name in names:
            columns = [row[2] for row in conn.execute("PRAGMA index_info(%s)" % name)]
            if columns and columns[0] == "app_registration_uuid":
                matching.append(name)
        self.assertTrue(matching)

    def test_pepper_synced_into_broker_config(self):
        broker, _ = self._boot()
        conn = self._connect()
        row = conn.execute(
            "SELECT value FROM broker_config WHERE key = 'BROKER_INSTANCE_PEPPER'"
        ).fetchone()
        self.assertIsNotNone(row)
        self.assertEqual(broker.config["BROKER_INSTANCE_PEPPER"], row[0])
        self.assertEqual(32, len(base64.b64decode(row[0])))

    def test_pepper_env_override_synced(self):
        self._boot({"BROKER_INSTANCE_PEPPER": "Zm9vYmFy"})
        conn = self._connect()
        row = conn.execute(
            "SELECT value FROM broker_config WHERE key = 'BROKER_INSTANCE_PEPPER'"
        ).fetchone()
        self.assertEqual("Zm9vYmFy", row[0])

    def test_second_boot_silent_and_idempotent(self):
        first, first_stdout = self._boot()
        second, second_stdout = self._boot()
        self.assertEqual(1, first_stdout.count(first.config["BROKER_ADMIN_PASS"]))
        self.assertEqual("", second_stdout)
        self.assertEqual(
            first.config["BROKER_INSTANCE_PEPPER"],
            second.config["BROKER_INSTANCE_PEPPER"],
        )
        self.assertEqual(
            first.config["BROKER_INSTANCE_PEPPER"],
            second.store.get_config("BROKER_INSTANCE_PEPPER"),
        )
        self.assertEqual(0o600, os.stat(self.db_path).st_mode & 0o777)

    def test_restart_preserves_rows(self):
        broker, _ = self._boot()
        broker.store.conn.execute(
            "INSERT INTO app_registrations (uuid, name) VALUES (?, ?)",
            ("0b6c3370-6d38-4b92-b9a2-bc2f1c9c1b0a", "proxmox-tool"),
        )
        broker.store.conn.commit()
        broker.store.close()
        conn = self._connect()
        row = conn.execute(
            "SELECT name FROM app_registrations WHERE uuid = ?",
            ("0b6c3370-6d38-4b92-b9a2-bc2f1c9c1b0a",),
        ).fetchone()
        self.assertEqual("proxmox-tool", row[0])


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