fix: Separate unit tests and fix asyncio in test_predict\n\n- Separated unit tests into unit and unit-minimal sessions to handle different dependency needs.\n- Added @pytest.mark.asyncio to TestModelServer.test_predict.

Author: chalmerlowe

noxfile.py

@@ -97,6 +97,7 @@
 # 'docfx' is excluded since it only needs to run in 'docs-presubmit'
 nox.options.sessions = [
     "unit",
+    "unit_minimal",
     "unit_ray",
     "unit_langchain",
     "unit_ag2",
@@ -219,6 +220,7 @@ def default(session):
         "--ignore=tests/unit/vertex_ag2",
         "--ignore=tests/unit/vertex_llama_index",
         "--ignore=tests/unit/architecture",
+        "--ignore=tests/unit/vertexai/test_generative_models.py",  # Exclude minimal tests
         os.path.join("tests", "unit"),
         *session.posargs,
     )
@@ -235,11 +237,7 @@ def default(session):
 
 @nox.session(python=UNIT_TEST_PYTHON_VERSIONS)
 def unit(session):
-    """Run the unit test suite."""
-    # First run the minimal GenAI tests
-    unit_genai_minimal_dependencies(session)
-
-    # Then run the default full test suite
+    """Run the default unit test suite."""
     default(session)
 
 
@@ -254,14 +252,18 @@ def unit_genai_minimal_dependencies(session):
     session.run(
         "py.test",
         "--quiet",
-        f"--junitxml=unit_{session.python}_sponge_log.xml",
-        # These tests require the PIL module
-        # "--ignore=TestGenerativeModels::test_image_mime_types",
+        f"--junitxml=unit_minimal_{session.python}_sponge_log.xml",
         os.path.join("tests", "unit", "vertexai", "test_generative_models.py"),
         *session.posargs,
     )
 
 
+@nox.session(python=UNIT_TEST_PYTHON_VERSIONS)
+def unit_minimal(session):
+    """Run the minimal GenAI unit tests."""
+    unit_genai_minimal_dependencies(session)
+
+
 @nox.session(python=["3.10", "3.11"])
 @nox.parametrize("ray", ["2.9.3", "2.33.0", "2.42.0", "2.47.1"])
 def unit_ray(session, ray):

tests/unit/aiplatform/test_prediction.py

@@ -3279,6 +3279,7 @@ def test_init_no_aip_predict_route(
 
         assert str(exception.value) == expected_message
 
+
     def test_health(self, model_server_env_mock, importlib_import_module_mock_twice):
         model_server = CprModelServer()
         client = TestClient(model_server.app)
@@ -3287,10 +3288,12 @@ def test_health(self, model_server_env_mock, importlib_import_module_mock_twice)
 
         assert response.status_code == 200
 
-    def test_predict(self, model_server_env_mock, importlib_import_module_mock_twice):
+    @pytest.mark.asyncio
+    async def test_predict(self, model_server_env_mock, importlib_import_module_mock_twice):
         model_server = CprModelServer()
         client = TestClient(model_server.app)
 
+
         with mock.patch.object(model_server.handler, "handle") as handle_mock:
             future = asyncio.Future()
             future.set_result(Response())