fixit: Update to Gen App Builder Tests

discoveryengine/documents_sample_test.py

@@ -28,40 +28,34 @@
 bigquery_table = "import_documents_test"
 
 
-def test_import_documents_gcs(capsys):
-    import_documents_sample.import_documents_sample(
+def test_import_documents_gcs():
+    operation_name = import_documents_sample.import_documents_sample(
         project_id=project_id,
         location=location,
         search_engine_id=search_engine_id,
         gcs_uri=gcs_uri,
     )
 
-    out, _ = capsys.readouterr()
+    assert "operations/import-documents" in operation_name
 
-    assert "operations/import-documents" in out
 
-
-def test_import_documents_bigquery(capsys):
-    import_documents_sample.import_documents_sample(
+def test_import_documents_bigquery():
+    operation_name = import_documents_sample.import_documents_sample(
         project_id=project_id,
         location=location,
         search_engine_id=search_engine_id,
         bigquery_dataset=bigquery_dataset,
         bigquery_table=bigquery_table,
     )
 
-    out, _ = capsys.readouterr()
-
-    assert "operations/import-documents" in out
+    assert "operations/import-documents" in operation_name
 
 
-def test_list_documents(capsys):
-    list_documents_sample.list_documents_sample(
+def test_list_documents():
+    response = list_documents_sample.list_documents_sample(
         project_id=project_id,
         location=location,
         search_engine_id=search_engine_id,
     )
 
-    out, _ = capsys.readouterr()
-
-    assert f"Documents in {search_engine_id}" in out
+    assert response

discoveryengine/get_operation_sample.py

@@ -14,7 +14,7 @@
 #
 
 # [START genappbuilder_get_operation]
-from google.cloud import discoveryengine_v1beta as genappbuilder
+from google.cloud import discoveryengine
 
 from google.longrunning import operations_pb2
 
@@ -23,9 +23,9 @@
 # operation_name = "YOUR_OPERATION_NAME"
 
 
-def get_operation_sample(operation_name: str) -> None:
+def get_operation_sample(operation_name: str) -> operations_pb2.Operation:
     # Create a client
-    client = genappbuilder.DocumentServiceClient()
+    client = discoveryengine.DocumentServiceClient()
 
     # Make GetOperation request
     request = operations_pb2.GetOperationRequest(name=operation_name)
@@ -34,5 +34,7 @@ def get_operation_sample(operation_name: str) -> None:
     # Print the Operation Information
     print(operation)
 
+    return operation
+
 
 # [END genappbuilder_get_operation]

discoveryengine/import_documents_sample.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 
-from google.cloud import discoveryengine_v1beta as genappbuilder
+from google.cloud import discoveryengine
 
 # TODO(developer): Uncomment these variables before running the sample.
 # project_id = "YOUR_PROJECT_ID"
@@ -38,9 +38,9 @@ def import_documents_sample(
     gcs_uri: str | None = None,
     bigquery_dataset: str | None = None,
     bigquery_table: str | None = None,
-) -> None:
+) -> str:
     # Create a client
-    client = genappbuilder.DocumentServiceClient()
+    client = discoveryengine.DocumentServiceClient()
 
     # The full resource name of the search engine branch.
     # e.g. projects/{project}/locations/{location}/dataStores/{data_store}/branches/{branch}
@@ -52,25 +52,25 @@ def import_documents_sample(
     )
 
     if gcs_uri:
-        request = genappbuilder.ImportDocumentsRequest(
+        request = discoveryengine.ImportDocumentsRequest(
             parent=parent,
-            gcs_source=genappbuilder.GcsSource(
+            gcs_source=discoveryengine.GcsSource(
                 input_uris=[gcs_uri], data_schema="custom"
             ),
             # Options: `FULL`, `INCREMENTAL`
-            reconciliation_mode=genappbuilder.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
+            reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
         )
     else:
-        request = genappbuilder.ImportDocumentsRequest(
+        request = discoveryengine.ImportDocumentsRequest(
             parent=parent,
-            bigquery_source=genappbuilder.BigQuerySource(
+            bigquery_source=discoveryengine.BigQuerySource(
                 project_id=project_id,
                 dataset_id=bigquery_dataset,
                 table_id=bigquery_table,
                 data_schema="custom",
             ),
             # Options: `FULL`, `INCREMENTAL`
-            reconciliation_mode=genappbuilder.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
+            reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
         )
 
     # Make the request
@@ -81,11 +81,13 @@ def import_documents_sample(
 
     # Once the operation is complete,
     # get information from operation metadata
-    metadata = genappbuilder.ImportDocumentsMetadata(operation.metadata)
+    metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)
 
     # Handle the response
     print(response)
     print(metadata)
 
+    return operation.operation.name
+
 
 # [END genappbuilder_import_documents]

discoveryengine/list_documents_sample.py

@@ -13,21 +13,19 @@
 # limitations under the License.
 #
 # [START genappbuilder_list_documents]
-from google.cloud import discoveryengine_v1beta as genappbuilder
+from typing import Any
+
+from google.cloud import discoveryengine
 
 # TODO(developer): Uncomment these variables before running the sample.
 # project_id = "YOUR_PROJECT_ID"
 # location = "YOUR_LOCATION" # Values: "global"
 # search_engine_id = "YOUR_SEARCH_ENGINE_ID"
 
 
-def list_documents_sample(
-    project_id: str,
-    location: str,
-    search_engine_id: str
-) -> None:
+def list_documents_sample(project_id: str, location: str, search_engine_id: str) -> Any:
     # Create a client
-    client = genappbuilder.DocumentServiceClient()
+    client = discoveryengine.DocumentServiceClient()
 
     # The full resource name of the search engine branch.
     # e.g. projects/{project}/locations/{location}/dataStores/{data_store}/branches/{branch}
@@ -44,5 +42,7 @@ def list_documents_sample(
     for result in response:
         print(result)
 
+    return response
+
 
 # [END genappbuilder_list_documents]

discoveryengine/list_operations_sample.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 
-from google.cloud import discoveryengine_v1beta as genappbuilder
+from google.cloud import discoveryengine
 
 from google.longrunning import operations_pb2
 
@@ -34,9 +34,9 @@ def list_operations_sample(
     location: str,
     search_engine_id: str,
     operations_filter: str | None = None,
-) -> None:
+) -> operations_pb2.ListOperationsResponse:
     # Create a client
-    client = genappbuilder.DocumentServiceClient()
+    client = discoveryengine.DocumentServiceClient()
 
     # The full resource name of the search engine branch.
     name = f"projects/{project_id}/locations/{location}/collections/default_collection/dataStores/{search_engine_id}"
@@ -54,5 +54,7 @@ def list_operations_sample(
     for operation in response.operations:
         print(operation)
 
+    return response
+
 
 # [END genappbuilder_list_operations]

discoveryengine/operations_sample_test.py

@@ -28,35 +28,34 @@
 operation_name = f"projects/{project_id}/locations/{location}/collections/default_collection/dataStores/{search_engine_id}/branches/0/operations/{operation_id}"
 
 
-def test_get_operation(capsys):
+def test_get_operation():
     try:
-        get_operation_sample.get_operation_sample(operation_name=operation_name)
+        operation = get_operation_sample.get_operation_sample(
+            operation_name=operation_name
+        )
+        assert operation
     except NotFound as e:
         print(e.message)
+        pass
 
-    out, _ = capsys.readouterr()
 
-    assert operation_id in out
-
-
-def test_list_operations(capsys):
-    list_operations_sample.list_operations_sample(
+def test_list_operations():
+    response = list_operations_sample.list_operations_sample(
         project_id=project_id,
         location=location,
         search_engine_id=search_engine_id,
     )
 
-    out, _ = capsys.readouterr()
-
-    assert operation_id in out
+    assert response
+    assert response.operations
 
 
-def test_poll_operation(capsys):
+def test_poll_operation():
     try:
-        poll_operation_sample.poll_operation_sample(operation_name=operation_name)
+        operation = poll_operation_sample.poll_operation_sample(
+            operation_name=operation_name
+        )
+        assert operation
     except NotFound as e:
         print(e.message)
-
-    out, _ = capsys.readouterr()
-
-    assert operation_id in out
+        pass

discoveryengine/poll_operation_sample.py

@@ -16,7 +16,7 @@
 # [START genappbuilder_poll_operation]
 from time import sleep
 
-from google.cloud import discoveryengine_v1beta as genappbuilder
+from google.cloud import discoveryengine
 
 from google.longrunning import operations_pb2
 
@@ -25,9 +25,11 @@
 # operation_name = "YOUR_OPERATION_NAME"
 
 
-def poll_operation_sample(operation_name: str, limit: int = 10):
+def poll_operation_sample(
+    operation_name: str, limit: int = 10
+) -> operations_pb2.Operation:
     # Create a client
-    client = genappbuilder.DocumentServiceClient()
+    client = discoveryengine.DocumentServiceClient()
 
     # Make GetOperation request
     request = operations_pb2.GetOperationRequest(name=operation_name)
@@ -44,5 +46,7 @@ def poll_operation_sample(operation_name: str, limit: int = 10):
         # Wait 10 seconds before polling again
         sleep(10)
 
+    return operation
+
 
 # [END genappbuilder_poll_operation]

discoveryengine/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-discoveryengine==0.7.0
+google-cloud-discoveryengine==0.8.0
 google-api-core==2.11.0

discoveryengine/search_sample.py

@@ -14,8 +14,9 @@
 #
 
 # [START genappbuilder_search]
+from typing import List
 
-from google.cloud import discoveryengine_v1beta as genappbuilder
+from google.cloud import discoveryengine
 
 # TODO(developer): Uncomment these variables before running the sample.
 # project_id = "YOUR_PROJECT_ID"
@@ -31,9 +32,9 @@ def search_sample(
     search_engine_id: str,
     serving_config_id: str,
     search_query: str,
-) -> None:
+) -> List[discoveryengine.SearchResponse.SearchResult]:
     # Create a client
-    client = genappbuilder.SearchServiceClient()
+    client = discoveryengine.SearchServiceClient()
 
     # The full resource name of the search engine serving config
     # e.g. projects/{project_id}/locations/{location}
@@ -44,13 +45,15 @@ def search_sample(
         serving_config=serving_config_id,
     )
 
-    request = genappbuilder.SearchRequest(
+    request = discoveryengine.SearchRequest(
         serving_config=serving_config,
         query=search_query,
     )
     response = client.search(request)
     for result in response.results:
         print(result)
 
+    return response.results
+
 
 # [END genappbuilder_search]

discoveryengine/search_sample_test.py

@@ -24,15 +24,13 @@
 search_query = "Google"
 
 
-def test_search(capsys):
-    search_sample.search_sample(
+def test_search():
+    response = search_sample.search_sample(
         project_id=project_id,
         location=location,
         search_engine_id=search_engine_id,
         serving_config_id=serving_config_id,
         search_query=search_query,
     )
 
-    out, _ = capsys.readouterr()
-
-    assert search_query in out
+    assert response