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()