test_app_startup.py

"""Exercise real Gradio launch with fake RabbitMQ and forecasting dependencies.

Run with: python -m unittest test_app_startup
Requires gradio[mcp] (5.42.0, as pinned by the Space) and aio-pika.
"""
import asyncio
from pathlib import Path
import runpy
import socket
import sys
import types
import unittest
from unittest.mock import AsyncMock, Mock, patch

import gradio as gr
from gradio.routes import App
from starlette.testclient import TestClient


class AppStartupTests(unittest.TestCase):
    def setUp(self):
        self.base = Mock()
        self.connected = False

        async def connect():
            self.connected = True

        async def close():
            self.connected = False

        self.base.connect = AsyncMock(side_effect=connect)
        self.base.close = AsyncMock(side_effect=close)
        self.base.is_connected = lambda: self.connected
        self.queue = Mock()
        self.queue.consume = AsyncMock(return_value="test-consumer")
        self.base.declare_queue_bind = AsyncMock(return_value=self.queue)
        servers = Mock(handle_chat_create=AsyncMock(), handle_images_generate=AsyncMock())
        self.modules = {}
        for name, attrs in {
            "config": {"settings": types.SimpleNamespace(
                RABBIT_EXCHANGE_TYPE="direct", RABBIT_ROUTING_KEY="execute.api",
                RABBIT_INSTANCE_NAME="startup-test")},
            "rabbit_base": {"RabbitBase": Mock(return_value=self.base)},
            "rabbit_repo": {"RabbitRepo": Mock()},
            "oa_server": {"OpenAIServers": Mock(return_value=servers)},
            "timesfm_backend": {"TimesFMBackend": Mock(), "StubImagesBackend": Mock()},
        }.items():
            module = types.ModuleType(name)
            module.__dict__.update(attrs)
            self.modules[name] = module
        self.modules["spaces"] = None

    def load_app(self, name="startup_test"):
        with patch.dict(sys.modules, self.modules):
            return runpy.run_path(str(Path(__file__).with_name("app.py")), run_name=name)

    def test_launch_starts_consumers_without_browser_and_shutdown_closes_connection(self):
        original_launch = gr.Blocks.launch
        launched = []

        def launch(blocks, **kwargs):
            # Execute the production entry point and real launch implementation.
            # Only change networking/blocking options so this test can finish.
            launched.append(blocks)
            with socket.socket() as sock:
                sock.bind(("127.0.0.1", 0))
                port = sock.getsockname()[1]
            kwargs.update(server_name="127.0.0.1", server_port=port,
                          prevent_thread_lock=True, debug=False, quiet=True,
                          inbrowser=False, share=False)
            return original_launch(blocks, **kwargs)

        try:
            with patch.object(gr.Blocks, "launch", launch):
                app = self.load_app("__main__")
            self.base.connect.assert_awaited_once()
            self.assertEqual(self.queue.consume.await_count, 2)
            self.assertTrue(app["demo"].mcp_server)
            for _ in range(2):
                self.assertIn("ready", asyncio.run(app["service_status"]()))
                self.assertEqual(asyncio.run(app["ping"]()), "ok")
            self.assertEqual(self.queue.consume.await_count, 2)
        finally:
            for blocks in launched:
                blocks.close(verbose=False)
        self.base.close.assert_awaited_once()
        self.assertFalse(self.connected)

    def test_startup_failure_propagates_and_closes_partial_connection(self):
        app = self.load_app()
        self.base.declare_queue_bind.side_effect = RuntimeError("binding failed")
        # Use Gradio's real application factory and lifespan composition.
        server = App.create_app(app["demo"],
                                app_kwargs={"lifespan": app["service_lifespan"]},
                                mcp_server=True)
        entered = False
        with self.assertRaises(Exception):
            with TestClient(server):
                entered = True
        self.assertFalse(entered, "Server became available despite failed startup")
        self.base.close.assert_awaited_once()
        self.assertFalse(self.connected)
        self.assertEqual(asyncio.run(app["service_status"]()), "ERROR: binding failed")


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