"""Tests for openbb_cli.codegen.router_gen — router module emission.""" from __future__ import annotations import ast from dataclasses import dataclass from openbb_cli.codegen import ( namespace_tree as nt, router_gen as rg, ) @dataclass class _Fetcher: module_name: str fetcher_class: str model_name: str is_streaming: bool = False @dataclass class _PostCommand: module_name: str function_name: str # --- _safe_segment / _module_name_for / _model_name_for --- def test_safe_segment_handles_special_chars_and_empty(): assert rg._safe_segment("Equity Price") == "equity_price" assert rg._safe_segment("___") == "ns" assert rg._safe_segment("Fixed-Income") == "fixed_income" def test_module_name_for_joins_dotted_path(): assert rg._module_name_for("equity") == "equity" assert rg._module_name_for("equity.price") == "equity_price" assert rg._module_name_for("equity.price.historical") == "equity_price_historical" def test_model_name_for_camelcases_dotted_path(): assert rg._model_name_for("equity.price.historical") == "EquityPriceHistorical" assert rg._model_name_for("equity.search") == "EquitySearch" assert rg._model_name_for("commodity.psd_report") == "CommodityPsdReport" # --- _format_docstring --- def test_format_docstring_single_line_appends_period(): out = rg._format_docstring("Get historical data") assert out == ' """Get historical data."""' def test_format_docstring_blank_text_renders_period_only(): out = rg._format_docstring(" ") assert out == ' """."""' def test_format_docstring_keeps_existing_terminator(): out = rg._format_docstring("What is X?") assert out == ' """What is X?"""' def test_format_docstring_multiline_uses_block_layout(): out = rg._format_docstring("First line.\nSecond line.\n\nThird line.") assert out.startswith(' """First line.') assert out.endswith(' """') # Continuation lines are indented to four spaces; blank lines stay blank. assert "\n Second line." in out assert "\n Third line." in out # --- generate_routers (full pipeline) --- def test_generate_routers_emits_top_router_with_get_command(): tree = nt.build_namespace_tree( {"equity.search": {"providers": ["fmp"], "description": "Search equities"}} ) fetchers = { "equity.search": _Fetcher( module_name="equity_search", fetcher_class="EquitySearchFetcher", model_name="EquitySearch", ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="fmp", fetchers_by_command=fetchers, post_commands_by_command={}, ) assert isinstance(out, rg.GeneratedRouters) assert len(out.routers) == 1 router = out.routers[0] assert router.module_name == "equity" assert router.entry_point_name == "equity" src = router.source assert 'router = Router(prefix="")' in src assert '@router.command(model="EquitySearch")' in src assert "async def search(" in src assert '"""Search equities."""' in src # Module parses ast.parse(src) def test_generate_routers_emits_stream_command_for_streaming_fetcher(): tree = nt.build_namespace_tree( {"crypto.feed": {"providers": ["wargame"], "description": "Live feed"}} ) fetchers = { "crypto.feed": _Fetcher( module_name="crypto_feed", fetcher_class="CryptoFeedFetcher", model_name="CryptoFeed", is_streaming=True, ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="wargame", fetchers_by_command=fetchers, post_commands_by_command={}, ) src = out.routers[0].source assert "from openbb_core.app.model.stream import OBBStream" in src assert '@router.command(model="CryptoFeed")' in src assert "OBBStream" in src ast.parse(src) def test_generate_routers_nests_sub_routers_via_include_router(): tree = nt.build_namespace_tree( { "equity.price.historical": { "providers": ["fmp"], "description": "OHLCV.", } } ) fetchers = { "equity.price.historical": _Fetcher( module_name="equity_price_historical", fetcher_class="EquityPriceHistoricalFetcher", model_name="EquityPriceHistorical", ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="fmp", fetchers_by_command=fetchers, post_commands_by_command={}, ) by_module = {r.module_name: r for r in out.routers} # Both top-level (equity) and nested (equity_price) modules emitted assert {"equity", "equity_price"} <= set(by_module) # Only the top-level router gets an entry point assert by_module["equity"].entry_point_name == "equity" assert by_module["equity_price"].entry_point_name is None # Top-level router includes the nested one with the correct prefix assert ( "from openbb_codegen.routers.equity_price import router as _price_router" in by_module["equity"].source ) assert ( 'router.include_router(_price_router, prefix="/price")' in by_module["equity"].source ) # The nested router defines the actual command assert "async def historical(" in by_module["equity_price"].source def test_generate_routers_emits_post_command_re_decoration(): tree = nt.build_namespace_tree( {"econometrics.regression": {"providers": [], "description": "Linear fit."}} ) posts = { "econometrics.regression": _PostCommand( module_name="econometrics_regression", function_name="regression", ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="tools", fetchers_by_command={}, post_commands_by_command=posts, ) src = out.routers[0].source assert ( "from openbb_codegen.providers.tools.models.econometrics_regression " "import regression as _regression" ) in src assert 'router.command(methods=["POST"])(_regression)' in src ast.parse(src) def test_generate_routers_falls_back_to_model_name_when_description_missing(): tree = nt.build_namespace_tree( {"equity.search": {"providers": ["fmp"]}}, # no description ) fetchers = { "equity.search": _Fetcher( module_name="equity_search", fetcher_class="EquitySearchFetcher", model_name="EquitySearch", ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="fmp", fetchers_by_command=fetchers, post_commands_by_command={}, ) assert '"""EquitySearch command."""' in out.routers[0].source def test_generate_routers_skips_command_without_matching_fetcher_or_post(): tree = nt.build_namespace_tree( {"equity.search": {"providers": ["fmp"], "description": "Search."}} ) out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="fmp", fetchers_by_command={}, # no fetcher registered post_commands_by_command={}, # no post command either ) src = out.routers[0].source # No command body emitted; only the router declaration remains assert "async def" not in src assert "router.command" not in src ast.parse(src) # --- top-level leaf commands land on the root router --- def test_generate_routers_top_level_command_becomes_root_command(): tree = nt.build_namespace_tree( {"eod": {"providers": ["eodhd"], "description": "End of day."}} ) fetchers = { "eod": _Fetcher( module_name="eod", fetcher_class="EodFetcher", model_name="Eod", ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="eodhd", fetchers_by_command=fetchers, post_commands_by_command={}, ) # No dedicated router module for a pure top-level command assert out.routers == [] assert len(out.root_commands) == 1 block = out.root_commands[0] assert '@router.command(model="Eod")' in block assert "async def eod(" in block assert out.root_has_stream is False def test_generate_routers_top_level_stream_command_sets_flag(): tree = nt.build_namespace_tree( {"feed": {"providers": ["wargame"], "description": "Live feed."}} ) fetchers = { "feed": _Fetcher( module_name="feed", fetcher_class="FeedFetcher", model_name="Feed", is_streaming=True, ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="wargame", fetchers_by_command=fetchers, post_commands_by_command={}, ) assert out.root_has_stream is True assert "OBBStream" in out.root_commands[0] def test_generate_routers_top_level_post_command_records_import(): tree = nt.build_namespace_tree( {"regression": {"providers": [], "description": "Linear fit."}} ) posts = { "regression": _PostCommand( module_name="regression", function_name="regression", ) } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="tools", fetchers_by_command={}, post_commands_by_command=posts, ) assert out.root_post_imports == [ ("openbb_codegen.providers.tools.models.regression", "regression") ] assert 'router.command(methods=["POST"])(_regression)' in out.root_commands[0] def test_generate_routers_top_level_command_without_fetcher_is_dropped(): tree = nt.build_namespace_tree( {"orphan": {"providers": ["fmp"], "description": "No fetcher."}} ) out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="fmp", fetchers_by_command={}, post_commands_by_command={}, ) assert out.routers == [] assert out.root_commands == [] def test_generate_routers_mixes_namespaces_and_root_commands(): tree = nt.build_namespace_tree( { "eod": {"providers": ["eodhd"], "description": "End of day."}, "calendar.ipos": {"providers": ["eodhd"], "description": "IPOs."}, } ) fetchers = { "eod": _Fetcher( module_name="eod", fetcher_class="EodFetcher", model_name="Eod" ), "calendar.ipos": _Fetcher( module_name="calendar_ipos", fetcher_class="CalendarIposFetcher", model_name="CalendarIpos", ), } out = rg.generate_routers( tree, package_name="openbb_codegen", provider_name="eodhd", fetchers_by_command=fetchers, post_commands_by_command={}, ) assert [r.module_name for r in out.routers] == ["calendar"] assert "async def ipos(" in out.routers[0].source assert len(out.root_commands) == 1 assert "async def eod(" in out.root_commands[0]