Refactor, modularize

Author: mliu-cloudera

llm-service/app/tests/provider_mocks/bedrock.py

@@ -36,6 +36,7 @@
 #  DATA.
 #
 import itertools
+from contextlib import AbstractContextManager
 from typing import Generator
 from unittest.mock import patch
 from urllib.parse import urljoin
@@ -47,9 +48,9 @@
 
 from app.config import settings
 from app.services.caii.types import ModelResponse
-from app.services.models import ModelProvider
 from app.services.models.providers import BedrockModelProvider
 from app.services.models.providers._model_provider import get_all_env_var_names
+from .utils import patch_env_vars
 
 TEXT_MODELS = [
     ("test.unavailable-text-model-v1", "NOT_AVAILABLE"),
@@ -64,14 +65,7 @@
 ]
 
 
-@pytest.fixture
-def mock_bedrock(monkeypatch) -> Generator[None, None, None]:
-    for name in BedrockModelProvider.get_env_var_names():
-        monkeypatch.setenv(name, "test")
-    for name in get_all_env_var_names() - BedrockModelProvider.get_env_var_names():
-        monkeypatch.delenv(name, raising=False)
-
-    # mock calls made directly through `requests`
+def _patch_requests() -> AbstractContextManager:
     bedrock_url_base = f"https://bedrock.{settings.aws_default_region}.amazonaws.com/"
     r_mock = responses.RequestsMock(assert_all_requests_are_fired=False)
     for model_id, availability in TEXT_MODELS + EMBEDDING_MODELS:
@@ -92,7 +86,10 @@ def mock_bedrock(monkeypatch) -> Generator[None, None, None]:
             },
         )
 
-    # mock calls made through `boto3`
+    return r_mock
+
+
+def _patch_boto3() -> AbstractContextManager:
     make_api_call = botocore.client.BaseClient._make_api_call
 
     def mock_make_api_call(self, operation_name: str, api_params: dict[str, str]):
@@ -151,29 +148,28 @@ def mock_make_api_call(self, operation_name: str, api_params: dict[str, str]):
             # passthrough
             return make_api_call(self, operation_name, api_params)
 
-    # mock reranking models, which are hard-coded in our app
-    def list_reranking_models() -> list[ModelResponse]:
-        return [
-            ModelResponse(model_id=model_id, name=model_id.upper())
-            for model_id, _ in RERANKING_MODELS
-        ]
-
-    with (
-        r_mock,
-        patch(
-            "botocore.client.BaseClient._make_api_call",
-            new=mock_make_api_call,
-        ),
-        patch(
-            "app.services.models.providers.BedrockModelProvider.list_reranking_models",
-            new=list_reranking_models,
-        ),
-        patch(  # work around a llama-index filter we have in list_llm_models()
-            "app.services.models.providers.bedrock.BEDROCK_MODELS",
-            new=BEDROCK_MODELS | {model_id: 128000 for model_id, _ in TEXT_MODELS},
-        ),
-    ):
-        yield
+    return patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
+
+
+@pytest.fixture
+def mock_bedrock(monkeypatch) -> Generator[None, None, None]:
+    with patch_env_vars(BedrockModelProvider):
+        with (
+            _patch_requests(),
+            _patch_boto3(),
+            patch(  # mock reranking models, which are hard-coded in our app
+                "app.services.models.providers.BedrockModelProvider.list_reranking_models",
+                new=lambda: [
+                    ModelResponse(model_id=model_id, name=model_id.upper())
+                    for model_id, _ in RERANKING_MODELS
+                ],
+            ),
+            patch(  # work around a llama-index filter we have in list_llm_models()
+                "app.services.models.providers.bedrock.BEDROCK_MODELS",
+                new=BEDROCK_MODELS | {model_id: 128000 for model_id, _ in TEXT_MODELS},
+            ),
+        ):
+            yield
 
 
 # TODO: move this test function to a discoverable place

llm-service/app/tests/provider_mocks/utils.py

@@ -0,0 +1,62 @@
+#
+#  CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
+#  (C) Cloudera, Inc. 2025
+#  All rights reserved.
+#
+#  Applicable Open Source License: Apache 2.0
+#
+#  NOTE: Cloudera open source products are modular software products
+#  made up of hundreds of individual components, each of which was
+#  individually copyrighted.  Each Cloudera open source product is a
+#  collective work under U.S. Copyright Law. Your license to use the
+#  collective work is as provided in your written agreement with
+#  Cloudera.  Used apart from the collective work, this file is
+#  licensed for your use pursuant to the open source license
+#  identified above.
+#
+#  This code is provided to you pursuant a written agreement with
+#  (i) Cloudera, Inc. or (ii) a third-party authorized to distribute
+#  this code. If you do not have a written agreement with Cloudera nor
+#  with an authorized and properly licensed third party, you do not
+#  have any rights to access nor to use this code.
+#
+#  Absent a written agreement with Cloudera, Inc. ("Cloudera") to the
+#  contrary, A) CLOUDERA PROVIDES THIS CODE TO YOU WITHOUT WARRANTIES OF ANY
+#  KIND; (B) CLOUDERA DISCLAIMS ANY AND ALL EXPRESS AND IMPLIED
+#  WARRANTIES WITH RESPECT TO THIS CODE, INCLUDING BUT NOT LIMITED TO
+#  IMPLIED WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY AND
+#  FITNESS FOR A PARTICULAR PURPOSE; (C) CLOUDERA IS NOT LIABLE TO YOU,
+#  AND WILL NOT DEFEND, INDEMNIFY, NOR HOLD YOU HARMLESS FOR ANY CLAIMS
+#  ARISING FROM OR RELATED TO THE CODE; AND (D)WITH RESPECT TO YOUR EXERCISE
+#  OF ANY RIGHTS GRANTED TO YOU FOR THE CODE, CLOUDERA IS NOT LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, PUNITIVE OR
+#  CONSEQUENTIAL DAMAGES INCLUDING, BUT NOT LIMITED TO, DAMAGES
+#  RELATED TO LOST REVENUE, LOST PROFITS, LOSS OF INCOME, LOSS OF
+#  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
+#  DATA.
+#
+import contextlib
+from typing import Iterator
+
+from pytest import MonkeyPatch
+
+from app.config import MODEL_PROVIDER_ENV_VAR_NAME
+from app.services.models.providers._model_provider import (
+    _ModelProvider,
+    get_all_env_var_names,
+)
+
+
+@contextlib.contextmanager
+def patch_env_vars(ModelProviderSubcls: type[_ModelProvider]) -> Iterator[None]:
+    """Set and unset environment variables for the given model provider."""
+    with MonkeyPatch.context() as monkeypatch:
+        for name in get_all_env_var_names():
+            monkeypatch.delenv(name, raising=False)
+        for name in ModelProviderSubcls.get_env_var_names():
+            monkeypatch.setenv(name, "test")
+        monkeypatch.setenv(
+            MODEL_PROVIDER_ENV_VAR_NAME,
+            ModelProviderSubcls.get_model_source(),
+        )
+        yield

llm-service/app/tests/services/test_models.py

@@ -35,37 +35,27 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
+from typing import Iterator
+
 import pytest
 
-from app.config import MODEL_PROVIDER_ENV_VAR_NAME
 from app.services import models
 from app.services.caii import caii
 from app.services.caii.types import ListEndpointEntry
 from app.services.models.providers import BedrockModelProvider
-from app.services.models.providers._model_provider import (
-    _ModelProvider,
-    get_all_env_var_names,
-)
+from app.services.models.providers._model_provider import _ModelProvider
+from app.tests.provider_mocks.utils import patch_env_vars
 
 
 @pytest.fixture(params=_ModelProvider.__subclasses__())
 def EnabledModelProvider(
     request: pytest.FixtureRequest,
     monkeypatch: pytest.MonkeyPatch,
-) -> type[_ModelProvider]:
-    """Sets and unsets environment variables for the given model provider."""
+) -> Iterator[type[_ModelProvider]]:
+    """Parametrize a test to run for each supported model provider."""
     ModelProviderSubcls: type[_ModelProvider] = request.param
-
-    for name in get_all_env_var_names():
-        monkeypatch.delenv(name, raising=False)
-    for name in ModelProviderSubcls.get_env_var_names():
-        monkeypatch.setenv(name, "test")
-    monkeypatch.setenv(
-        MODEL_PROVIDER_ENV_VAR_NAME,
-        ModelProviderSubcls.get_model_source(),
-    )
-
-    return ModelProviderSubcls
+    with patch_env_vars(ModelProviderSubcls):
+        yield ModelProviderSubcls
 
 
 class TestListAvailableModels: