diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 3a6a955f807..04c582f6013 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -770,3 +770,105 @@ def test_apply_default_settings(): mock_add_checks.assert_called_once_with([]) + + +def test_initialize_core_endpoints(): + """ + Test that _initialize_core_endpoints correctly sets up all core router endpoints. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "anthropic/test-model", + "api_key": "fake-api-key", + }, + } + ] + ) + + router._initialize_core_endpoints() + + core_endpoints = [ + "amoderation", + "aanthropic_messages", + "agenerate_content", + "aadapter_generate_content", + "aresponses", + "afile_delete", + "afile_content", + "responses", + "aget_responses", + "acancel_responses", + "adelete_responses", + "alist_input_items", + "_arealtime", + "acreate_fine_tuning_job", + "acancel_fine_tuning_job", + "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job", + "afile_list", + "aimage_edit", + "allm_passthrough_route", + ] + + for endpoint in core_endpoints: + assert hasattr(router, endpoint) + assert callable(getattr(router, endpoint)) + + +def test_initialize_specialized_endpoints(): + """ + Test that _initialize_specialized_endpoints correctly sets up specialized endpoints. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/test-model", + "api_key": "fake-api-key", + }, + } + ] + ) + + router._initialize_specialized_endpoints() + + specialized_endpoints = [ + "avector_store_search", + "avector_store_create", + "vector_store_search", + "vector_store_create", + "agenerate_content", + "generate_content", + "agenerate_content_stream", + "generate_content_stream", + "aocr", + "ocr", + "asearch", + "search", + "avideo_generation", + "video_generation", + "avideo_list", + "video_list", + "avideo_status", + "video_status", + "avideo_content", + "video_content", + "avideo_remix", + "video_remix", + "acreate_container", + "create_container", + "alist_containers", + "list_containers", + "aretrieve_container", + "retrieve_container", + "adelete_container", + "delete_container", + ] + + for endpoint in specialized_endpoints: + assert hasattr(router, endpoint) + assert callable(getattr(router, endpoint))