chore: Combine reasoning engine system tests into one test method.

PiperOrigin-RevId: 632336753

tests/system/vertexai/test_reasoning_engines.py

@@ -18,20 +18,13 @@
 import pytest
 from google import auth
 from google.api_core import exceptions
+from google.cloud import storage
 import vertexai
 from tests.system.aiplatform import e2e_base
 from vertexai.preview import reasoning_engines
 
 
-class CapitalizeEngine:
-    """A sample Reasoning Engine."""
-
-    def set_up(self):
-        pass
-
-    def query(self, input: str) -> str:
-        """Capitalizes the input."""
-        return input.upper()
+_BLOB_FILENAME = vertexai.reasoning_engines._reasoning_engines._BLOB_FILENAME
 
 
 @pytest.mark.usefixtures(
@@ -53,81 +46,41 @@ def test_langchain_template(self, shared_state):
             staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
             credentials=credentials,
         )
+        # Test prebuilt langchain_template
         created_app = reasoning_engines.ReasoningEngine.create(
             reasoning_engines.LangchainAgent(model="gemini-1.0-pro"),
             requirements=["google-cloud-aiplatform[reasoningengine,langchain]"],
+            display_name="test-display-name",
+            description="test-description",
+            gcs_dir_name="test-gcs-dir-name",
         )
         shared_state.setdefault("resources", [])
         shared_state["resources"].append(created_app)  # Deletion at teardown.
+        got_app = reasoning_engines.ReasoningEngine(created_app.resource_name)
+
+        # Test resource attributes
+        assert isinstance(created_app.resource_name, str)
+        assert got_app.resource_name == created_app.resource_name
+        assert got_app.gca_resource.name == got_app.resource_name
+        assert got_app.gca_resource.display_name == "test-display-name"
+        assert got_app.gca_resource.description == "test-description"
+
+        # Test operation schemas
+        assert got_app.operation_schemas() == created_app.operation_schemas()
+
+        # Test query response
+        # (Wrap in a try-except block because of non-determinism from Gemini.)
         try:
             response = created_app.query(input="hello")
             assert response.get("input") == "hello"
-            assert isinstance(created_app.resource_name, str)
-            got_app = reasoning_engines.ReasoningEngine(created_app.resource_name)
-            assert got_app.resource_name == created_app.resource_name
-            assert got_app.operation_schemas() == created_app.operation_schemas()
             response = got_app.query(input="hello")
             assert response.get("input") == "hello"
         except exceptions.FailedPrecondition as e:
             print(e)
 
-    def test_create_reasoning_engine_gcs_dir_name(self, shared_state):
-        # https://github.com/googleapis/python-aiplatform/issues/3650
-        super().setup_method()
-        credentials, _ = auth.default(
-            scopes=["https://www.googleapis.com/auth/cloud-platform"]
-        )
-        vertexai.init(
-            project=e2e_base._PROJECT,
-            location=e2e_base._LOCATION,
-            staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
-            credentials=credentials,
-        )
-        created_app = reasoning_engines.ReasoningEngine.create(
-            reasoning_engine=CapitalizeEngine(),
-            gcs_dir_name="test-gcs-dir-name",
-        )
-        shared_state.setdefault("resources", [])
-        shared_state["resources"].append(created_app)  # Deletion at teardown.
-        assert created_app.query(input="hello") == "HELLO"
-
-    def test_create_reasoning_engine_resource_attributes(self, shared_state):
-        super().setup_method()
-        credentials, _ = auth.default(
-            scopes=["https://www.googleapis.com/auth/cloud-platform"]
-        )
-        vertexai.init(
-            project=e2e_base._PROJECT,
-            location=e2e_base._LOCATION,
-            staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
-            credentials=credentials,
-        )
-        created_app = reasoning_engines.ReasoningEngine.create(
-            reasoning_engine=CapitalizeEngine(),
-            reasoning_engine_name="test-reasoning-engine-name",
-            display_name="test-display-name",
-            description="test-description",
-        )
-        shared_state.setdefault("resources", [])
-        shared_state["resources"].append(created_app)  # Deletion at teardown.
-        assert created_app.gca_resource.name == "test-reasoning-engine-name"
-        assert created_app.gca_resource.display_name == "test-display-name"
-        assert created_app.gca_resource.description == "test-description"
-
-    def test_create_reasoning_engine_operation_schemas(self, shared_state):
-        super().setup_method()
-        credentials, _ = auth.default(
-            scopes=["https://www.googleapis.com/auth/cloud-platform"]
-        )
-        vertexai.init(
-            project=e2e_base._PROJECT,
-            location=e2e_base._LOCATION,
-            staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
-            credentials=credentials,
-        )
-        created_app = reasoning_engines.ReasoningEngine.create(
-            reasoning_engine=CapitalizeEngine(),
-        )
-        shared_state.setdefault("resources", [])
-        shared_state["resources"].append(created_app)  # Deletion at teardown.
-        assert created_app.operation_schemas() == []
+        # Test GCS Bucket subdirectory creation
+        # Original: https://github.com/googleapis/python-aiplatform/issues/3650
+        client = storage.Client(project=e2e_base._PROJECT)
+        bucket = client.bucket(shared_state["staging_bucket_name"])
+        assert bucket.exists()
+        assert bucket.get_blob(f"test-gcs-dir-name/{_BLOB_FILENAME}").exists()