data-collector-mcp/tests/test_db.py

from __future__ import annotations

import unittest
from unittest.mock import MagicMock, patch

from data_collector_mcp import db


class DatabaseTests(unittest.TestCase):
    def test_database_url_requires_value(self) -> None:
        with patch.dict("os.environ", {}, clear=True):
            with self.assertRaisesRegex(ValueError, "DATABASE_URL is required"):
                db.database_url()

    def test_database_url_rejects_sqlite(self) -> None:
        with patch.dict("os.environ", {"DATABASE_URL": "sqlite:///test.db"}, clear=True):
            with self.assertRaisesRegex(ValueError, "SQLite is not supported"):
                db.database_url()

    def test_database_url_accepts_postgres(self) -> None:
        url = "postgresql+psycopg2://user:pass@localhost:5432/data-collector-mcp"
        with patch.dict("os.environ", {"DATABASE_URL": url}, clear=True):
            self.assertEqual(db.database_url(), url)

    def test_check_database_connection_executes_select(self) -> None:
        connection = MagicMock()
        connect_context = MagicMock()
        connect_context.__enter__.return_value = connection
        engine = MagicMock()
        engine.connect.return_value = connect_context

        with patch.object(db, "sql_engine", return_value=engine):
            db.check_database_connection()

        connection.execute.assert_called_once()

    def test_check_database_connection_propagates_errors(self) -> None:
        engine = MagicMock()
        engine.connect.side_effect = RuntimeError("connection failed")

        with patch.object(db, "sql_engine", return_value=engine):
            with self.assertRaisesRegex(RuntimeError, "connection failed"):
                db.check_database_connection()


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