Merge pull request #205 from IBM/isready_test

db_isready.py unit tests

Author: crivetimihai

tests/unit/mcpgateway/cache/test_session_registry.py

@@ -767,9 +767,7 @@ async def test_shutdown_with_redis_error(monkeypatch):
     """shutdown() should swallow Redis / PubSub aclose() errors."""
 
     # Tell the registry that the Redis extras are present
-    monkeypatch.setattr(
-        "mcpgateway.cache.session_registry.REDIS_AVAILABLE", True
-    )
+    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)
 
     # ── fake PubSub object ────────────────────────────────────────────────
     mock_pubsub = AsyncMock(name="MockPubSub")
@@ -789,7 +787,7 @@ async def test_shutdown_with_redis_error(monkeypatch):
             backend="redis",
             redis_url="redis://localhost:6379",
         )
-        await registry.initialize()     # calls mock_redis.pubsub()
+        await registry.initialize()  # calls mock_redis.pubsub()
 
         # must swallow both aclose() exceptions
         await registry.shutdown()

tests/unit/mcpgateway/test_db_isready.py

@@ -0,0 +1,219 @@
+# -*- coding: utf-8 -*-
+# Standard
+import asyncio
+import sys
+
+# First-Party
+# Import the module under test after patching where necessary
+import mcpgateway.utils.db_isready as db_isready
+
+# Third-Party
+import pytest
+
+# ---------------------------------------------------------------------------
+# Helper test doubles
+# ---------------------------------------------------------------------------
+
+
+class _DummyConn:
+    """A no‑op DBAPI connection that always succeeds on ``execute``."""
+
+    def execute(self, _):
+        return 1  # pragma: no cover
+
+    # Context‑manager support -------------------------------------------------
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+class _DummyEngine:
+    """Mimics the minimal SQLAlchemy *Engine* interface needed by db_isready."""
+
+    def __init__(self, succeed_after: int = 1):
+        self._attempts = 0
+        self._succeed_after = max(1, succeed_after)
+
+    def connect(self):
+        # Import inside the method so SQLAlchemy is only required when tests run
+        # Third-Party
+        from sqlalchemy.exc import OperationalError  # pylint: disable=C0415
+
+        self._attempts += 1
+        if self._attempts < self._succeed_after:
+            raise OperationalError("SELECT 1", {}, Exception("boom"))
+        return _DummyConn()
+
+    # Expose attempts for assertions
+    @property
+    def attempts(self):  # noqa: D401 – simple accessor
+        return self._attempts
+
+
+# ---------------------------------------------------------------------------
+# Unit‑tests – utilities first, then public API
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "raw",
+    [
+        "postgresql://alice:secret@db/mydb",
+        "error password=reallys3cret param=value",
+    ],
+)
+def test_sanitize_masks_sensitive_parts(raw):
+    """Anything that looks like credentials must be replaced by ***."""
+
+    redacted = db_isready._sanitize(raw)
+
+    # The replacement text must contain at least one asterisk block signalling masking
+    assert "***" in redacted
+
+    # And **no** piece of the original secret text may survive
+    assert "secret" not in redacted
+    assert "reallys3cret" not in redacted
+
+
+@pytest.mark.parametrize(
+    "url, expected",
+    [
+        ("sqlite:///:memory:", ":memory:"),  # SQLAlchemy represents memory DB with literal string
+        ("postgresql://u:[email protected]:5432/mcp", "db.example.com:5432/mcp"),
+    ],
+)
+def test_format_target_variants(url, expected):
+    """_format_target should create concise human readable targets."""
+
+    assert db_isready._format_target(db_isready.make_url(url)) == expected
+
+
+def test_wait_for_db_ready_success(monkeypatch):
+    """A healthy database should succeed on the first attempt."""
+
+    dummy = _DummyEngine(succeed_after=1)
+
+    def _fake_create_engine(_url, **kwargs):
+        _fake_create_engine.kwargs = kwargs  # type: ignore[attr-defined]
+        return dummy
+
+    monkeypatch.setattr(db_isready, "create_engine", _fake_create_engine)
+    monkeypatch.setattr(db_isready.time, "sleep", lambda *_: None)
+
+    db_isready.wait_for_db_ready(
+        database_url="postgresql://user:pw@localhost:5432/mcp",
+        max_tries=3,
+        interval=0.001,
+        timeout=1,
+        sync=True,
+    )
+
+    assert dummy.attempts == 1
+    assert _fake_create_engine.kwargs["connect_args"]["connect_timeout"] == 1  # type: ignore[attr-defined]
+
+
+def test_wait_for_db_ready_retries_then_succeeds(monkeypatch):
+    """OperationalError should trigger retries until the connection works."""
+
+    dummy = _DummyEngine(succeed_after=3)
+    monkeypatch.setattr(db_isready, "create_engine", lambda *_a, **_k: dummy)
+    monkeypatch.setattr(db_isready.time, "sleep", lambda *_: None)
+
+    db_isready.wait_for_db_ready(
+        database_url="postgresql://u:p@db/mcp",
+        max_tries=5,
+        interval=0.0001,
+        timeout=2,
+        sync=True,
+    )
+
+    assert dummy.attempts == 3
+
+
+def test_wait_for_db_ready_exhausts_and_raises(monkeypatch):
+    """After *max_tries* failures the helper must raise RuntimeError."""
+
+    dummy = _DummyEngine(succeed_after=999)
+    monkeypatch.setattr(db_isready, "create_engine", lambda *_a, **_k: dummy)
+    monkeypatch.setattr(db_isready.time, "sleep", lambda *_: None)
+
+    with pytest.raises(RuntimeError, match="Database not ready after 3 attempts"):
+        db_isready.wait_for_db_ready(
+            database_url="sqlite:///tmp.db",
+            max_tries=3,
+            interval=0.001,
+            timeout=1,
+            sync=True,
+        )
+    assert dummy.attempts == 3
+
+
+def test_wait_for_db_ready_invalid_parameters():
+    """Zero or negative timing parameters are rejected immediately."""
+
+    with pytest.raises(RuntimeError):
+        db_isready.wait_for_db_ready(max_tries=0)
+    with pytest.raises(RuntimeError):
+        db_isready.wait_for_db_ready(interval=0)
+    with pytest.raises(RuntimeError):
+        db_isready.wait_for_db_ready(timeout=0)
+
+
+def test_wait_for_db_ready_async_path(monkeypatch):
+    """Async path should off‑load probe into executor without blocking."""
+
+    dummy = _DummyEngine(succeed_after=1)
+    monkeypatch.setattr(db_isready, "create_engine", lambda *_a, **_k: dummy)
+    monkeypatch.setattr(db_isready.time, "sleep", lambda *_: None)
+
+    # Create a dedicated loop so we can patch run_in_executor cleanly
+    loop = asyncio.new_event_loop()
+
+    async def _fake_run_in_executor(_executor, func, *args):  # noqa: D401
+        # Execute the probe synchronously (no thread) then return dummy future
+        func(*args)
+        fut = loop.create_future()
+        fut.set_result(None)
+        return fut
+
+    loop.run_in_executor = _fake_run_in_executor  # type: ignore[assignment]
+    monkeypatch.setattr(asyncio, "get_event_loop", lambda: loop)
+
+    db_isready.wait_for_db_ready(
+        database_url="postgresql://u:p@db/mcp",
+        max_tries=2,
+        interval=0.001,
+        timeout=1,
+        sync=False,
+    )
+
+    assert dummy.attempts == 1
+    loop.close()
+
+
+def test_parse_cli_roundtrip(monkeypatch):
+    """All CLI flags should be parsed into the expected Namespace values."""
+
+    argv = [
+        "db_isready.py",
+        "--database-url",
+        "postgresql://u:p@db/mcp",
+        "--max-tries",
+        "7",
+        "--interval",
+        "0.5",
+        "--timeout",
+        "3",
+        "--log-level",
+        "DEBUG",
+    ]
+    monkeypatch.setattr(sys, "argv", argv)
+
+    ns = db_isready._parse_cli()
+    assert ns.database_url == "postgresql://u:p@db/mcp"
+    assert ns.max_tries == 7
+    assert ns.interval == 0.5
+    assert ns.timeout == 3
+    assert ns.log_level == "DEBUG"