diff --git a/discoveryengine/documents_sample_test.py b/discoveryengine/documents_sample_test.py index c9bfb82fff0..153bb8e0544 100644 --- a/discoveryengine/documents_sample_test.py +++ b/discoveryengine/documents_sample_test.py @@ -28,21 +28,19 @@ 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, @@ -50,18 +48,14 @@ def test_import_documents_bigquery(capsys): 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 diff --git a/discoveryengine/get_operation_sample.py b/discoveryengine/get_operation_sample.py index 3e5e899c64b..40d12f591f3 100644 --- a/discoveryengine/get_operation_sample.py +++ b/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] diff --git a/discoveryengine/import_documents_sample.py b/discoveryengine/import_documents_sample.py index 72d19082b12..21b564f3fcf 100644 --- a/discoveryengine/import_documents_sample.py +++ b/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] diff --git a/discoveryengine/list_documents_sample.py b/discoveryengine/list_documents_sample.py index d61e594beba..e9c4b1ee025 100644 --- a/discoveryengine/list_documents_sample.py +++ b/discoveryengine/list_documents_sample.py @@ -13,7 +13,9 @@ # 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" @@ -21,13 +23,9 @@ # 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] diff --git a/discoveryengine/list_operations_sample.py b/discoveryengine/list_operations_sample.py index f7993cf569f..f95503caaf7 100644 --- a/discoveryengine/list_operations_sample.py +++ b/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] diff --git a/discoveryengine/operations_sample_test.py b/discoveryengine/operations_sample_test.py index 92b46844813..7534e518a4a 100644 --- a/discoveryengine/operations_sample_test.py +++ b/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 diff --git a/discoveryengine/poll_operation_sample.py b/discoveryengine/poll_operation_sample.py index 13e7eb435d2..30bba6e6462 100644 --- a/discoveryengine/poll_operation_sample.py +++ b/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] diff --git a/discoveryengine/requirements.txt b/discoveryengine/requirements.txt index b280e5b9702..8a6056b4daa 100644 --- a/discoveryengine/requirements.txt +++ b/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 diff --git a/discoveryengine/search_sample.py b/discoveryengine/search_sample.py index afe1510f4b3..fa2d0a87e82 100644 --- a/discoveryengine/search_sample.py +++ b/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,7 +45,7 @@ def search_sample( serving_config=serving_config_id, ) - request = genappbuilder.SearchRequest( + request = discoveryengine.SearchRequest( serving_config=serving_config, query=search_query, ) @@ -52,5 +53,7 @@ def search_sample( for result in response.results: print(result) + return response.results + # [END genappbuilder_search] diff --git a/discoveryengine/search_sample_test.py b/discoveryengine/search_sample_test.py index 076da2138e9..eb7ef04365e 100644 --- a/discoveryengine/search_sample_test.py +++ b/discoveryengine/search_sample_test.py @@ -24,8 +24,8 @@ 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, @@ -33,6 +33,4 @@ def test_search(capsys): search_query=search_query, ) - out, _ = capsys.readouterr() - - assert search_query in out + assert response