ria-ran--sim-drop/tests/test_database.py

import pytest
from datetime import datetime, timedelta

from server.database import (
    add_available_subscriber,
    add_subscriber,
    check_rate_limit,
    claim_available,
    count_available,
    export_all,
    get_connection,
    get_subscriber_by_id,
    get_subscriber_by_imsi,
    init_db,
    list_unsynced,
    mark_synced,
)


@pytest.fixture
def db_conn(tmp_path):
    db_path = str(tmp_path / "test.db")
    init_db(db_path)
    conn = get_connection(db_path)
    yield conn
    conn.close()


def _sub(**overrides):
    defaults = dict(
        imsi="302720000000001",
        iccid="8910302720000000001",
        ki="aa" * 16,
        opc="bb" * 16,
        mac_address="aa:bb:cc:dd:ee:ff",
    )
    defaults.update(overrides)
    return defaults


class TestCRUD:
    def test_add_and_retrieve_by_imsi(self, db_conn):
        d = _sub()
        sid = add_subscriber(db_conn, **d)
        sub = get_subscriber_by_imsi(db_conn, d["imsi"])
        assert sub is not None
        assert sub.imsi == d["imsi"]
        assert sub.id == sid

    def test_add_and_retrieve_by_id(self, db_conn):
        d = _sub()
        sid = add_subscriber(db_conn, **d)
        sub = get_subscriber_by_id(db_conn, sid)
        assert sub is not None
        assert sub.iccid == d["iccid"]

    def test_duplicate_imsi_raises(self, db_conn):
        add_subscriber(db_conn, **_sub())
        with pytest.raises(ValueError, match="Duplicate"):
            add_subscriber(db_conn, **_sub())

    def test_duplicate_iccid_raises(self, db_conn):
        add_subscriber(db_conn, **_sub())
        with pytest.raises(ValueError):
            add_subscriber(db_conn, **_sub(imsi="302720000000002"))

    def test_not_found_returns_none(self, db_conn):
        assert get_subscriber_by_imsi(db_conn, "999999999999999") is None
        assert get_subscriber_by_id(db_conn, 99999) is None

    def test_synced_to_core_defaults_false(self, db_conn):
        add_subscriber(db_conn, **_sub())
        sub = get_subscriber_by_imsi(db_conn, _sub()["imsi"])
        assert sub.synced_to_core is False


class TestRateLimiting:
    def test_first_request_allowed(self, db_conn):
        assert check_rate_limit(db_conn, "aa:bb:cc:dd:ee:ff", 1) is True

    def test_second_request_blocked(self, db_conn):
        add_subscriber(db_conn, **_sub())
        assert check_rate_limit(db_conn, "aa:bb:cc:dd:ee:ff", 1) is False

    def test_different_mac_allowed(self, db_conn):
        add_subscriber(db_conn, **_sub())
        assert check_rate_limit(db_conn, "11:22:33:44:55:66", 1) is True

    def test_expired_window_allows_again(self, db_conn):
        add_subscriber(db_conn, **_sub())
        past = (datetime.utcnow() - timedelta(hours=2)).isoformat()
        db_conn.execute("UPDATE subscribers SET created_at = ?", (past,))
        db_conn.commit()
        assert check_rate_limit(db_conn, "aa:bb:cc:dd:ee:ff", 1) is True


class TestSync:
    def test_list_unsynced_after_insert(self, db_conn):
        add_subscriber(db_conn, **_sub())
        assert len(list_unsynced(db_conn)) == 1

    def test_mark_synced_removes_from_unsynced(self, db_conn):
        add_subscriber(db_conn, **_sub())
        mark_synced(db_conn, _sub()["imsi"])
        assert len(list_unsynced(db_conn)) == 0

    def test_synced_flag_set(self, db_conn):
        add_subscriber(db_conn, **_sub())
        mark_synced(db_conn, _sub()["imsi"])
        sub = get_subscriber_by_imsi(db_conn, _sub()["imsi"])
        assert sub.synced_to_core is True


class TestBatchClaim:
    def test_available_defaults_and_count(self, db_conn):
        assert count_available(db_conn) == 0
        add_available_subscriber(db_conn, "302720000000010", "8910302720000000010",
                                 "aa" * 16, "bb" * 16)
        assert count_available(db_conn) == 1
        sub = get_subscriber_by_imsi(db_conn, "302720000000010")
        assert sub.status == "available"
        assert sub.mac_address is None

    def test_on_demand_add_is_claimed(self, db_conn):
        add_subscriber(db_conn, **_sub())
        sub = get_subscriber_by_imsi(db_conn, _sub()["imsi"])
        assert sub.status == "claimed"

    def test_claim_assigns_and_decrements(self, db_conn):
        add_available_subscriber(db_conn, "302720000000010", "8910302720000000010",
                                 "aa" * 16, "bb" * 16)
        claimed = claim_available(db_conn, "aa:bb:cc:dd:ee:ff")
        assert claimed is not None
        assert claimed.status == "claimed"
        assert claimed.mac_address == "aa:bb:cc:dd:ee:ff"
        assert count_available(db_conn) == 0

    def test_claim_returns_none_when_empty(self, db_conn):
        assert claim_available(db_conn, "aa:bb:cc:dd:ee:ff") is None

    def test_claim_is_fifo_and_unique(self, db_conn):
        for i in range(3):
            add_available_subscriber(db_conn, f"30272000000002{i}",
                                     f"891030272000000002{i}", "aa" * 16, "bb" * 16)
        first = claim_available(db_conn, "11:11:11:11:11:11")
        second = claim_available(db_conn, "22:22:22:22:22:22")
        assert first.imsi != second.imsi
        assert first.id < second.id  # oldest first
        assert count_available(db_conn) == 1


class TestMigration:
    def test_adds_status_column_to_legacy_db(self, tmp_path):
        import sqlite3
        db_path = str(tmp_path / "legacy.db")
        legacy = sqlite3.connect(db_path)
        legacy.executescript(
            "CREATE TABLE subscribers (id INTEGER PRIMARY KEY AUTOINCREMENT, "
            "imsi TEXT UNIQUE NOT NULL, iccid TEXT UNIQUE NOT NULL, ki TEXT NOT NULL, "
            "opc TEXT NOT NULL, mac_address TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, "
            "synced_to_core BOOLEAN DEFAULT 0);"
        )
        legacy.execute(
            "INSERT INTO subscribers (imsi, iccid, ki, opc) VALUES "
            "('302720000000099', '8910302720000000099', 'aa', 'bb')"
        )
        legacy.commit()
        legacy.close()

        init_db(db_path)  # should add the status column, defaulting old rows to 'claimed'
        conn = get_connection(db_path)
        try:
            sub = get_subscriber_by_imsi(conn, "302720000000099")
            assert sub.status == "claimed"
        finally:
            conn.close()


class TestExport:
    def test_export_returns_list_with_correct_keys(self, db_conn):
        add_subscriber(db_conn, **_sub())
        rows = export_all(db_conn)
        assert len(rows) == 1
        for key in ("imsi", "iccid", "ki", "opc", "mac_address", "created_at", "synced_to_core"):
            assert key in rows[0]

    def test_export_empty_db(self, db_conn):
        assert export_all(db_conn) == []