From bd88cbf13ced867d9aa75f4a32ba013846c62772 Mon Sep 17 00:00:00 2001 From: Mariatta Wijaya Date: Mon, 8 May 2023 22:16:30 -0700 Subject: [PATCH 1/8] Feat: Parametrize the tests for multidb support Remove "database" argument from Query and AggregationQuery constructors. Use the "database" from the client instead. Once set in the client, the "database" will be used throughout and cannot be re-set. Parametrize the tests where-ever clients are used. Use the `system-tests-named-db` in the system test. --- google/cloud/datastore/_http.py | 51 +- google/cloud/datastore/aggregation.py | 46 +- google/cloud/datastore/batch.py | 28 +- google/cloud/datastore/client.py | 67 +-- google/cloud/datastore/constants.py | 4 - google/cloud/datastore/helpers.py | 12 +- google/cloud/datastore/key.py | 40 +- google/cloud/datastore/query.py | 62 +-- google/cloud/datastore/transaction.py | 20 +- google/cloud/datastore_v1/types/entity.py | 6 +- noxfile.py | 6 +- tests/system/_helpers.py | 2 +- tests/system/conftest.py | 16 +- tests/system/index.yaml | 16 +- tests/system/test_aggregation_query.py | 39 +- tests/system/test_allocate_reserve_ids.py | 15 +- tests/system/test_put.py | 28 +- tests/system/test_query.py | 57 +- tests/system/test_read_consistency.py | 10 +- tests/system/test_transaction.py | 12 +- tests/system/utils/populate_datastore.py | 4 +- tests/unit/test__http.py | 234 +++++--- tests/unit/test_aggregation.py | 80 +-- tests/unit/test_batch.py | 275 +++++----- tests/unit/test_client.py | 627 ++++++++++++---------- tests/unit/test_helpers.py | 37 +- tests/unit/test_key.py | 182 ++++--- tests/unit/test_query.py | 433 ++++++++------- tests/unit/test_transaction.py | 237 ++++---- 29 files changed, 1488 insertions(+), 1158 deletions(-) delete mode 100644 google/cloud/datastore/constants.py diff --git a/google/cloud/datastore/_http.py b/google/cloud/datastore/_http.py index 61209e98..065ee15b 100644 --- a/google/cloud/datastore/_http.py +++ b/google/cloud/datastore/_http.py @@ -61,6 +61,7 @@ def _request( client_info, retry=None, timeout=None, + database=None, ): """Make a request over the Http transport to the Cloud Datastore API. @@ -90,6 +91,9 @@ def _request( :type timeout: float or tuple(float, float) :param timeout: (Optional) timeout for the request + :type database: str + :param database: The database to make the request for. + :rtype: str :returns: The string response content from the API call. :raises: :class:`google.cloud.exceptions.GoogleCloudError` if the @@ -101,6 +105,7 @@ def _request( "User-Agent": user_agent, connection_module.CLIENT_INFO_HEADER: user_agent, } + _update_headers(headers, project, database) api_url = build_api_url(project, method, base_url) requester = http.request @@ -138,6 +143,7 @@ def _rpc( response_pb_cls, retry=None, timeout=None, + database=None, ): """Make a protobuf RPC request. @@ -177,7 +183,14 @@ def _rpc( req_data = request_pb._pb.SerializeToString() kwargs = _make_retry_timeout_kwargs(retry, timeout) response = _request( - http, project, method, req_data, base_url, client_info, **kwargs + http, + project, + method, + req_data, + base_url, + client_info, + database=database, + **kwargs, ) return response_pb_cls.deserialize(response) @@ -237,6 +250,7 @@ def lookup(self, request, retry=None, timeout=None): request_pb = _make_request_pb(request, _datastore_pb2.LookupRequest) project_id = request_pb.project_id + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -247,6 +261,7 @@ def lookup(self, request, retry=None, timeout=None): _datastore_pb2.LookupResponse, retry=retry, timeout=timeout, + database=database, ) def run_query(self, request, retry=None, timeout=None): @@ -267,7 +282,7 @@ def run_query(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.RunQueryRequest) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -278,6 +293,7 @@ def run_query(self, request, retry=None, timeout=None): _datastore_pb2.RunQueryResponse, retry=retry, timeout=timeout, + database=database, ) def run_aggregation_query(self, request, retry=None, timeout=None): @@ -300,7 +316,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None): request, _datastore_pb2.RunAggregationQueryRequest ) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -311,6 +327,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None): _datastore_pb2.RunAggregationQueryResponse, retry=retry, timeout=timeout, + database=database, ) def begin_transaction(self, request, retry=None, timeout=None): @@ -331,7 +348,7 @@ def begin_transaction(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.BeginTransactionRequest) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -342,6 +359,7 @@ def begin_transaction(self, request, retry=None, timeout=None): _datastore_pb2.BeginTransactionResponse, retry=retry, timeout=timeout, + database=database, ) def commit(self, request, retry=None, timeout=None): @@ -362,7 +380,7 @@ def commit(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.CommitRequest) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -373,6 +391,7 @@ def commit(self, request, retry=None, timeout=None): _datastore_pb2.CommitResponse, retry=retry, timeout=timeout, + database=database, ) def rollback(self, request, retry=None, timeout=None): @@ -393,7 +412,7 @@ def rollback(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.RollbackRequest) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -404,6 +423,7 @@ def rollback(self, request, retry=None, timeout=None): _datastore_pb2.RollbackResponse, retry=retry, timeout=timeout, + database=database, ) def allocate_ids(self, request, retry=None, timeout=None): @@ -424,7 +444,7 @@ def allocate_ids(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.AllocateIdsRequest) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -435,6 +455,7 @@ def allocate_ids(self, request, retry=None, timeout=None): _datastore_pb2.AllocateIdsResponse, retry=retry, timeout=timeout, + database=database, ) def reserve_ids(self, request, retry=None, timeout=None): @@ -455,7 +476,7 @@ def reserve_ids(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.ReserveIdsRequest) project_id = request_pb.project_id - + database = request_pb.database_id return _rpc( self.client._http, project_id, @@ -466,4 +487,18 @@ def reserve_ids(self, request, retry=None, timeout=None): _datastore_pb2.ReserveIdsResponse, retry=retry, timeout=timeout, + database=database, ) + + +def _update_headers(headers, project_id, database_id=None): + """Update the request headers. + + Pass the project id, or optionally the database_id if provided. + """ + if database_id: + headers[ + "x-goog-request-params" + ] = f"project_id={project_id}&database_id={database_id}" + else: + headers["x-goog-request-params"] = f"project_id={project_id}" diff --git a/google/cloud/datastore/aggregation.py b/google/cloud/datastore/aggregation.py index 7632e215..25fbf992 100644 --- a/google/cloud/datastore/aggregation.py +++ b/google/cloud/datastore/aggregation.py @@ -22,7 +22,6 @@ from google.cloud.datastore_v1.types import query as query_pb2 from google.cloud.datastore import helpers from google.cloud.datastore.query import _pb_from_query -from google.cloud.datastore.constants import DEFAULT_DATABASE _NOT_FINISHED = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED @@ -124,18 +123,6 @@ def project(self): """ return self._nested_query._project or self._client.project - @property - def database(self): - """Get the database for this AggregationQuery. - :rtype: str - :returns: The database for the query. - """ - if self._nested_query._database or ( - self._nested_query._database == DEFAULT_DATABASE - ): - return self._nested_query._database - return self._client.database - @property def namespace(self): """The nested query's namespace @@ -389,7 +376,7 @@ def _next_page(self): partition_id = entity_pb2.PartitionId( project_id=self._aggregation_query.project, - database_id=self._aggregation_query.database, + database_id=self.client.database, namespace_id=self._aggregation_query.namespace, ) @@ -400,15 +387,15 @@ def _next_page(self): if self._timeout is not None: kwargs["timeout"] = self._timeout - + request = { + "project_id": self._aggregation_query.project, + "partition_id": partition_id, + "read_options": read_options, + "aggregation_query": query_pb, + } + helpers._set_database_id_to_request(request, self.client.database) response_pb = self.client._datastore_api.run_aggregation_query( - request={ - "project_id": self._aggregation_query.project, - "database_id": self._aggregation_query.database, - "partition_id": partition_id, - "read_options": read_options, - "aggregation_query": query_pb, - }, + request=request, **kwargs, ) @@ -421,14 +408,15 @@ def _next_page(self): query_pb = query_pb2.AggregationQuery() query_pb._pb.CopyFrom(old_query_pb._pb) # copy for testability + request = { + "project_id": self._aggregation_query.project, + "partition_id": partition_id, + "read_options": read_options, + "aggregation_query": query_pb, + } + helpers._set_database_id_to_request(request, self.client.database) response_pb = self.client._datastore_api.run_aggregation_query( - request={ - "project_id": self._aggregation_query.project, - "database_id": self._aggregation_query.database, - "partition_id": partition_id, - "read_options": read_options, - "aggregation_query": query_pb, - }, + request=request, **kwargs, ) diff --git a/google/cloud/datastore/batch.py b/google/cloud/datastore/batch.py index 93ff5c81..276c0822 100644 --- a/google/cloud/datastore/batch.py +++ b/google/cloud/datastore/batch.py @@ -23,7 +23,6 @@ from google.cloud.datastore import helpers from google.cloud.datastore_v1.types import datastore as _datastore_pb2 -from google.cloud.datastore.constants import DEFAULT_DATABASE class Batch(object): @@ -228,10 +227,7 @@ def put(self, entity): if self.project != entity.key.project: raise ValueError("Key must be from same project as batch") - entity_key_database = entity.key.database - if entity_key_database is None: - entity_key_database = DEFAULT_DATABASE - if self.database != entity_key_database: + if self.database != entity.key.database: raise ValueError("Key must be from same database as batch") if entity.key.is_partial: @@ -261,10 +257,7 @@ def delete(self, key): if self.project != key.project: raise ValueError("Key must be from same project as batch") - key_db = key.database - if key_db is None: - key_db = DEFAULT_DATABASE - if self.database != key_db: + if self.database != key.database: raise ValueError("Key must be from same database as batch") key_pb = key.to_protobuf() @@ -303,14 +296,17 @@ def _commit(self, retry, timeout): if timeout is not None: kwargs["timeout"] = timeout + request = { + "project_id": self.project, + "mode": mode, + "transaction": self._id, + "mutations": self._mutations, + } + + helpers._set_database_id_to_request(request, self._client.database) + commit_response_pb = self._client._datastore_api.commit( - request={ - "project_id": self.project, - "database_id": self.database, - "mode": mode, - "transaction": self._id, - "mutations": self._mutations, - }, + request=request, **kwargs, ) diff --git a/google/cloud/datastore/client.py b/google/cloud/datastore/client.py index 105b590d..8b3d89d4 100644 --- a/google/cloud/datastore/client.py +++ b/google/cloud/datastore/client.py @@ -25,7 +25,6 @@ from google.cloud.datastore import helpers from google.cloud.datastore._http import HTTPDatastoreAPI from google.cloud.datastore.batch import Batch -from google.cloud.datastore.constants import DEFAULT_DATABASE from google.cloud.datastore.entity import Entity from google.cloud.datastore.key import Key from google.cloud.datastore.query import Query @@ -127,7 +126,7 @@ def _extended_lookup( retry=None, timeout=None, read_time=None, - database=DEFAULT_DATABASE, + database=None, ): """Repeat lookup until all keys found (unless stop requested). @@ -204,13 +203,14 @@ def _extended_lookup( read_options = helpers.get_read_options(eventual, transaction_id, read_time) while loop_num < _MAX_LOOPS: # loop against possible deferred. loop_num += 1 + request = { + "project_id": project, + "keys": key_pbs, + "read_options": read_options, + } + helpers._set_database_id_to_request(request, database) lookup_response = datastore_api.lookup( - request={ - "project_id": project, - "database_id": database, - "keys": key_pbs, - "read_options": read_options, - }, + request=request, **kwargs, ) @@ -298,10 +298,9 @@ def __init__( credentials=None, client_info=_CLIENT_INFO, client_options=None, + database=None, _http=None, _use_grpc=None, - *, - database=DEFAULT_DATABASE, ): emulator_host = os.getenv(DATASTORE_EMULATOR_HOST) @@ -318,12 +317,12 @@ def __init__( client_options=client_options, _http=_http, ) - self.database = database self.namespace = namespace self._client_info = client_info self._client_options = client_options self._batch_stack = _LocalStack() self._datastore_api_internal = None + self._database = database if _use_grpc is None: self._use_grpc = _USE_GRPC @@ -358,6 +357,11 @@ def base_url(self, value): """Setter for API base URL.""" self._base_url = value + @property + def database(self): + """Getter for database""" + return self._database + @property def _datastore_api(self): """Getter for a wrapped API object.""" @@ -562,7 +566,6 @@ def get_multi( entity_pbs = _extended_lookup( datastore_api=self._datastore_api, project=self.project, - database=self.database, key_pbs=[key.to_protobuf() for key in keys], eventual=eventual, missing=missing, @@ -571,6 +574,7 @@ def get_multi( retry=retry, timeout=timeout, read_time=read_time, + database=self.database, ) if missing is not None: @@ -753,12 +757,13 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None): kwargs = _make_retry_timeout_kwargs(retry, timeout) + request = { + "project_id": incomplete_key.project, + "keys": incomplete_key_pbs, + } + helpers._set_database_id_to_request(request, self.database) response_pb = self._datastore_api.allocate_ids( - request={ - "project_id": incomplete_key.project, - "database_id": incomplete_key.database, - "keys": incomplete_key_pbs, - }, + request=request, **kwargs, ) allocated_ids = [ @@ -801,7 +806,7 @@ def transaction(self, **kwargs): def query(self, **kwargs): """Proxy to :class:`google.cloud.datastore.query.Query`. - Passes our ``project`` and our ``database``. + Passes our ``project``. Using query to search a datastore: @@ -855,10 +860,7 @@ def do_something_with(entity): raise TypeError("Cannot pass client") if "project" in kwargs: raise TypeError("Cannot pass project") - if "database" in kwargs: - raise TypeError("Cannot pass database") kwargs["project"] = self.project - kwargs["database"] = self.database if "namespace" not in kwargs: kwargs["namespace"] = self.namespace return Query(self, **kwargs) @@ -1000,12 +1002,13 @@ def reserve_ids_sequential(self, complete_key, num_ids, retry=None, timeout=None key_pbs.append(key.to_protobuf()) kwargs = _make_retry_timeout_kwargs(retry, timeout) + request = { + "project_id": complete_key.project, + "keys": key_pbs, + } + helpers._set_database_id_to_request(request, self.database) self._datastore_api.reserve_ids( - request={ - "project_id": complete_key.project, - "database_id": complete_key.database, - "keys": key_pbs, - }, + request=request, **kwargs, ) return None @@ -1052,12 +1055,14 @@ def reserve_ids_multi(self, complete_keys, retry=None, timeout=None): kwargs = _make_retry_timeout_kwargs(retry, timeout) key_pbs = [key.to_protobuf() for key in complete_keys] + request = { + "project_id": complete_keys[0].project, + "keys": key_pbs, + } + helpers._set_database_id_to_request(request, complete_keys[0].database) + self._datastore_api.reserve_ids( - request={ - "project_id": complete_keys[0].project, - "database_id": complete_keys[0].database, - "keys": key_pbs, - }, + request=request, **kwargs, ) diff --git a/google/cloud/datastore/constants.py b/google/cloud/datastore/constants.py deleted file mode 100644 index 29f083a6..00000000 --- a/google/cloud/datastore/constants.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Constants for Datastore.""" - -DEFAULT_DATABASE = "" -"""Datastore default database.""" diff --git a/google/cloud/datastore/helpers.py b/google/cloud/datastore/helpers.py index e09cb82e..fa6730ce 100644 --- a/google/cloud/datastore/helpers.py +++ b/google/cloud/datastore/helpers.py @@ -29,7 +29,6 @@ from google.cloud.datastore_v1.types import entity as entity_pb2 from google.cloud.datastore.entity import Entity from google.cloud.datastore.key import Key -from google.cloud.datastore.constants import DEFAULT_DATABASE from google.protobuf import timestamp_pb2 @@ -301,7 +300,8 @@ def key_from_protobuf(pb): project = None if pb.partition_id.project_id: # Simple field (string) project = pb.partition_id.project_id - database = DEFAULT_DATABASE + database = None + if pb.partition_id.database_id: # Simple field (string) database = pb.partition_id.database_id namespace = None @@ -490,6 +490,14 @@ def _set_protobuf_value(value_pb, val): setattr(value_pb, attr, val) +def _set_database_id_to_request(request, database_id=None): + """ + Set the "database_id" field to the request only if it was provided. + """ + if database_id is not None: + request["database_id"] = database_id + + class GeoPoint(object): """Simple container for a geo point value. diff --git a/google/cloud/datastore/key.py b/google/cloud/datastore/key.py index 3e6b856c..b0c8e5d6 100644 --- a/google/cloud/datastore/key.py +++ b/google/cloud/datastore/key.py @@ -21,7 +21,6 @@ from google.cloud._helpers import _to_bytes from google.cloud.datastore import _app_engine_key_pb2 -from google.cloud.datastore.constants import DEFAULT_DATABASE _DATABASE_ID_TEMPLATE = ( @@ -115,7 +114,8 @@ def __init__(self, *path_args, **kwargs): self._flat_path = path_args parent = self._parent = kwargs.get("parent") self._namespace = kwargs.get("namespace") - self._database = kwargs.get("database") or DEFAULT_DATABASE + self._database = kwargs.get("database") + project = kwargs.get("project") self._project = _validate_project(project, parent) # _flat_path, _parent, _database, _namespace, and _project must be set @@ -130,8 +130,7 @@ def __eq__(self, other): Completed keys compare equal if they have the same path, project, database, and namespace. - (Note that "" and None are considered the same in the specific case of - databases, as these both refer to the default database.) + (Note that database=None is considered to refer to the default database.) :rtype: bool :returns: True if the keys compare equal, else False. @@ -142,18 +141,11 @@ def __eq__(self, other): if self.is_partial or other.is_partial: return False - self_database = self.database - if self_database is None: # pragma: NO COVER - self_database = DEFAULT_DATABASE - other_database = other.database - if other_database is None: # pragma: NO COVER - other_database = DEFAULT_DATABASE - return ( self.flat_path == other.flat_path and self.project == other.project - and self_database == other_database and self.namespace == other.namespace + and self.database == other.database ) def __ne__(self, other): @@ -164,8 +156,7 @@ def __ne__(self, other): Completed keys compare equal if they have the same path, project, database, and namespace. - (Note that "" and None are considered the same in the specific case of - databases, as these both refer to the default database.) + (Note that database=None is considered to refer to the default database.) :rtype: bool :returns: False if the keys compare equal, else True. @@ -178,12 +169,10 @@ def __hash__(self): :rtype: int :returns: a hash of the key's state. """ - return ( - hash(self.flat_path) - + hash(self.project) - + hash(self.namespace) - + hash(self.database) - ) + hash_val = hash(self.flat_path) + hash(self.project) + hash(self.namespace) + if self.database: + hash_val + hash(self.database) + return hash_val @staticmethod def _parse_path(path_args): @@ -256,6 +245,8 @@ def _combine_args(self): self._namespace = self._parent.namespace if self._project is not None and self._project != self._parent.project: raise ValueError("Child project must agree with parent's.") + if self._database is not None and self._database != self._parent.database: + raise ValueError("Child database must agree with parent's.") self._database = self._parent.database self._project = self._parent.project @@ -316,7 +307,8 @@ def to_protobuf(self): """ key = _entity_pb2.Key() key.partition_id.project_id = self.project - key.partition_id.database_id = self.database or DEFAULT_DATABASE + if self.database: + key.partition_id.database_id = self.database if self.namespace: key.partition_id.namespace_id = self.namespace @@ -407,9 +399,7 @@ def from_legacy_urlsafe(cls, urlsafe): namespace = _get_empty(reference.name_space, "") _check_database_id(reference.database_id) flat_path = _get_flat_path(reference.path) - return cls( - *flat_path, project=project, database=DEFAULT_DATABASE, namespace=namespace - ) + return cls(*flat_path, project=project, namespace=namespace) @property def is_partial(self): @@ -551,7 +541,7 @@ def __repr__(self): equal to the default values. """ repr = "" diff --git a/google/cloud/datastore/query.py b/google/cloud/datastore/query.py index 5d0b807c..14335268 100644 --- a/google/cloud/datastore/query.py +++ b/google/cloud/datastore/query.py @@ -176,11 +176,6 @@ class Query(object): :type distinct_on: sequence of string :param distinct_on: field names used to group query results. - :type database: str - :param database: - (optional) The database associated with the query. If not passed, - uses the client's value. - :raises: ValueError if ``project`` is not passed and no implicit default is set. """ @@ -208,8 +203,6 @@ def __init__( projection=(), order=(), distinct_on=(), - *, - database=None, ): self._client = client @@ -222,15 +215,6 @@ def __init__( else: self._project = None - # database defaults to None to allow distinguishing between an explicit - # query against the default database (for which you would pass database="") - # and a fallback to the client (for which you would either simply omit the parameter, - # or explicitly pass database=None) - if database is None: - self._database = client.database - else: - self._database = database - if namespace: self._namespace = namespace elif hasattr(client, "namespace"): @@ -261,17 +245,6 @@ def project(self): """ return self._project or self._client.project - @property - def database(self): - """Get the database for this Query. - - :rtype: str or None - :returns: The database for the query. - """ - if self._database is not None: - return self._database - return self._client.database - @property def namespace(self): """This query's namespace @@ -817,7 +790,7 @@ def _next_page(self): partition_id = entity_pb2.PartitionId( project_id=self._query.project, - database_id=self._query.database, + database_id=self.client.database, namespace_id=self._query.namespace, ) @@ -829,14 +802,17 @@ def _next_page(self): if self._timeout is not None: kwargs["timeout"] = self._timeout + request = { + "project_id": self._query.project, + "partition_id": partition_id, + "read_options": read_options, + "query": query_pb, + } + + helpers._set_database_id_to_request(request, self.client.database) + response_pb = self.client._datastore_api.run_query( - request={ - "project_id": self._query.project, - "database_id": self._query.database, - "partition_id": partition_id, - "read_options": read_options, - "query": query_pb, - }, + request=request, **kwargs, ) @@ -854,14 +830,16 @@ def _next_page(self): query_pb.start_cursor = response_pb.batch.skipped_cursor query_pb.offset -= response_pb.batch.skipped_results + request = { + "project_id": self._query.project, + "partition_id": partition_id, + "read_options": read_options, + "query": query_pb, + } + helpers._set_database_id_to_request(request, self.client.database) + response_pb = self.client._datastore_api.run_query( - request={ - "project_id": self._query.project, - "database_id": self._query.database, - "partition_id": partition_id, - "read_options": read_options, - "query": query_pb, - }, + request=request, **kwargs, ) diff --git a/google/cloud/datastore/transaction.py b/google/cloud/datastore/transaction.py index fd1d580b..73a9c68e 100644 --- a/google/cloud/datastore/transaction.py +++ b/google/cloud/datastore/transaction.py @@ -18,6 +18,8 @@ from google.cloud.datastore_v1.types import TransactionOptions from google.protobuf import timestamp_pb2 +from google.cloud.datastore.helpers import _set_database_id_to_request + def _make_retry_timeout_kwargs(retry, timeout): """Helper: make optional retry / timeout kwargs dict.""" @@ -225,9 +227,10 @@ def begin(self, retry=None, timeout=None): request = { "project_id": self.project, - "database_id": self.database, "transaction_options": self._options, } + _set_database_id_to_request(request, self._client.database) + try: response_pb = self._client._datastore_api.begin_transaction( request=request, **kwargs @@ -259,14 +262,13 @@ def rollback(self, retry=None, timeout=None): try: # No need to use the response it contains nothing. - self._client._datastore_api.rollback( - request={ - "project_id": self.project, - "database_id": self.database, - "transaction": self._id, - }, - **kwargs - ) + request = { + "project_id": self.project, + "transaction": self._id, + } + + _set_database_id_to_request(request, self._client.database) + self._client._datastore_api.rollback(request=request, **kwargs) finally: super(Transaction, self).rollback() # Clear our own ID in case this gets accidentally reused. diff --git a/google/cloud/datastore_v1/types/entity.py b/google/cloud/datastore_v1/types/entity.py index 2b1b78d5..ed66e490 100644 --- a/google/cloud/datastore_v1/types/entity.py +++ b/google/cloud/datastore_v1/types/entity.py @@ -39,8 +39,7 @@ class PartitionId(proto.Message): r"""A partition ID identifies a grouping of entities. The grouping is always by project. database. and namespace, however the namespace ID may be - empty. Default ("") and empty database ID's both refer to the default - database and are considered equivalent. + empty. Empty database ID refers to the default database. A partition ID contains several dimensions: project ID, database ID, and namespace ID. @@ -64,8 +63,7 @@ class PartitionId(proto.Message): belong. database_id (str): If not empty, the ID of the database to which - the entities belong. Empty and "" both correspond - to the default database. + the entities belong. Empty corresponds to the default database. namespace_id (str): If not empty, the ID of the namespace to which the entities belong. diff --git a/noxfile.py b/noxfile.py index 2faa3578..84ae80a4 100644 --- a/noxfile.py +++ b/noxfile.py @@ -230,8 +230,7 @@ def install_systemtest_dependencies(session, *constraints): @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) @nox.parametrize("disable_grpc", [False, True]) -@nox.parametrize("use_named_db", [False, True]) -def system(session, disable_grpc, use_named_db): +def system(session, disable_grpc): """Run the system test suite.""" constraints_path = str( CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" @@ -245,8 +244,6 @@ def system(session, disable_grpc, use_named_db): # Install pyopenssl for mTLS testing. if os.environ.get("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false") == "true": session.install("pyopenssl") - if use_named_db and os.environ.get("RUN_NAMED_DB_TESTS", "false") == "false": - session.skip("RUN_NAMED_DB_TESTS is set to false, skipping") system_test_exists = os.path.exists(system_test_path) system_test_folder_exists = os.path.exists(system_test_folder_path) @@ -259,7 +256,6 @@ def system(session, disable_grpc, use_named_db): env = {} if disable_grpc: env["GOOGLE_CLOUD_DISABLE_GRPC"] = "True" - env["SYSTEM_TESTS_DATABASE"] = "system-tests-named-db" if use_named_db else "" # Run py.test against the system tests. if system_test_exists: diff --git a/tests/system/_helpers.py b/tests/system/_helpers.py index e35ca363..13735625 100644 --- a/tests/system/_helpers.py +++ b/tests/system/_helpers.py @@ -19,7 +19,7 @@ from test_utils.system import unique_resource_id _DATASTORE_DATABASE = "SYSTEM_TESTS_DATABASE" -TEST_DATABASE = os.getenv(_DATASTORE_DATABASE) +TEST_DATABASE = os.getenv(_DATASTORE_DATABASE, "system-tests-named-db") EMULATOR_DATASET = os.getenv(DATASTORE_DATASET) diff --git a/tests/system/conftest.py b/tests/system/conftest.py index d985c852..1840556b 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -24,28 +24,32 @@ def in_emulator(): return _helpers.EMULATOR_DATASET is not None +@pytest.fixture(scope="session") +def database_id(request): + return request.param + + @pytest.fixture(scope="session") def test_namespace(): return _helpers.unique_id("ns") @pytest.fixture(scope="session") -def datastore_client(test_namespace): - database = "" +def datastore_client(test_namespace, database_id): if _helpers.TEST_DATABASE is not None: - database = _helpers.TEST_DATABASE + database_id = _helpers.TEST_DATABASE if _helpers.EMULATOR_DATASET is not None: http = requests.Session() # Un-authorized. client = datastore.Client( project=_helpers.EMULATOR_DATASET, - database=database, + database=database_id, namespace=test_namespace, _http=http, ) else: - client = datastore.Client(database=database, namespace=test_namespace) + client = datastore.Client(database=database_id, namespace=test_namespace) - assert client.database == database + assert client.database == database_id return client diff --git a/tests/system/index.yaml b/tests/system/index.yaml index 08a50d09..f9cc2a5b 100644 --- a/tests/system/index.yaml +++ b/tests/system/index.yaml @@ -30,4 +30,18 @@ indexes: - kind: Character properties: - name: Character - - name: appearances \ No newline at end of file + - name: appearances + +- kind: Character + ancestor: yes + properties: + - name: alive + - name: family + - name: appearances + + +- kind: Character + ancestor: yes + properties: + - name: family + - name: appearances diff --git a/tests/system/test_aggregation_query.py b/tests/system/test_aggregation_query.py index b912e96b..51045003 100644 --- a/tests/system/test_aggregation_query.py +++ b/tests/system/test_aggregation_query.py @@ -40,8 +40,8 @@ def _do_fetch(aggregation_query, **kw): @pytest.fixture(scope="session") -def aggregation_query_client(datastore_client): - return _helpers.clone_client(datastore_client, namespace=None) +def aggregation_query_client(datastore_client, database_id=None): + return _helpers.clone_client(datastore_client, namespace=None, database=database_id) @pytest.fixture(scope="session") @@ -69,7 +69,8 @@ def nested_query(aggregation_query_client, ancestor_key): return _make_query(aggregation_query_client, ancestor_key) -def test_aggregation_query_default(aggregation_query_client, nested_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_aggregation_query_default(aggregation_query_client, nested_query, database_id): query = nested_query aggregation_query = aggregation_query_client.aggregation_query(query) @@ -81,7 +82,10 @@ def test_aggregation_query_default(aggregation_query_client, nested_query): assert r.value == 8 -def test_aggregation_query_with_alias(aggregation_query_client, nested_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_aggregation_query_with_alias( + aggregation_query_client, nested_query, database_id +): query = nested_query aggregation_query = aggregation_query_client.aggregation_query(query) @@ -93,7 +97,10 @@ def test_aggregation_query_with_alias(aggregation_query_client, nested_query): assert r.value > 0 -def test_aggregation_query_with_limit(aggregation_query_client, nested_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_aggregation_query_with_limit( + aggregation_query_client, nested_query, database_id +): query = nested_query aggregation_query = aggregation_query_client.aggregation_query(query) @@ -113,8 +120,9 @@ def test_aggregation_query_with_limit(aggregation_query_client, nested_query): assert r.value == 2 +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) def test_aggregation_query_multiple_aggregations( - aggregation_query_client, nested_query + aggregation_query_client, nested_query, database_id ): query = nested_query @@ -128,7 +136,10 @@ def test_aggregation_query_multiple_aggregations( assert r.value > 0 -def test_aggregation_query_add_aggregation(aggregation_query_client, nested_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_aggregation_query_add_aggregation( + aggregation_query_client, nested_query, database_id +): from google.cloud.datastore.aggregation import CountAggregation query = nested_query @@ -143,7 +154,10 @@ def test_aggregation_query_add_aggregation(aggregation_query_client, nested_quer assert r.value > 0 -def test_aggregation_query_add_aggregations(aggregation_query_client, nested_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_aggregation_query_add_aggregations( + aggregation_query_client, nested_query, database_id +): from google.cloud.datastore.aggregation import CountAggregation query = nested_query @@ -159,8 +173,9 @@ def test_aggregation_query_add_aggregations(aggregation_query_client, nested_que assert r.value > 0 +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) def test_aggregation_query_add_aggregations_duplicated_alias( - aggregation_query_client, nested_query + aggregation_query_client, nested_query, database_id ): from google.cloud.datastore.aggregation import CountAggregation from google.api_core.exceptions import BadRequest @@ -187,8 +202,9 @@ def test_aggregation_query_add_aggregations_duplicated_alias( _do_fetch(aggregation_query) +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) def test_aggregation_query_with_nested_query_filtered( - aggregation_query_client, nested_query + aggregation_query_client, nested_query, database_id ): query = nested_query @@ -210,8 +226,9 @@ def test_aggregation_query_with_nested_query_filtered( assert r.value == 6 +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) def test_aggregation_query_with_nested_query_multiple_filters( - aggregation_query_client, nested_query + aggregation_query_client, nested_query, database_id ): query = nested_query diff --git a/tests/system/test_allocate_reserve_ids.py b/tests/system/test_allocate_reserve_ids.py index f934d067..2d7c3700 100644 --- a/tests/system/test_allocate_reserve_ids.py +++ b/tests/system/test_allocate_reserve_ids.py @@ -12,10 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import warnings +from . import _helpers -def test_client_allocate_ids(datastore_client): + +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_allocate_ids(datastore_client, database_id): num_ids = 10 allocated_keys = datastore_client.allocate_ids( datastore_client.key("Kind"), @@ -32,7 +36,8 @@ def test_client_allocate_ids(datastore_client): assert len(unique_ids) == num_ids -def test_client_reserve_ids_sequential(datastore_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_reserve_ids_sequential(datastore_client, database_id): num_ids = 10 key = datastore_client.key("Kind", 1234) @@ -41,7 +46,8 @@ def test_client_reserve_ids_sequential(datastore_client): datastore_client.reserve_ids_sequential(key, num_ids) -def test_client_reserve_ids_deprecated(datastore_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_reserve_ids_deprecated(datastore_client, database_id): num_ids = 10 key = datastore_client.key("Kind", 1234) @@ -53,7 +59,8 @@ def test_client_reserve_ids_deprecated(datastore_client): assert "reserve_ids_sequential" in str(warned[0].message) -def test_client_reserve_ids_multi(datastore_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_reserve_ids_multi(datastore_client, database_id): key1 = datastore_client.key("Kind", 1234) key2 = datastore_client.key("Kind", 1235) diff --git a/tests/system/test_put.py b/tests/system/test_put.py index ae1ed057..4cb5f6e8 100644 --- a/tests/system/test_put.py +++ b/tests/system/test_put.py @@ -54,7 +54,8 @@ def _get_post(datastore_client, id_or_name=None, post_content=None): @pytest.mark.parametrize( "name,key_id", [(None, None), ("post1", None), (None, 123456789)] ) -def test_client_put(datastore_client, entities_to_delete, name, key_id): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_put(datastore_client, entities_to_delete, name, key_id, database_id): entity = _get_post(datastore_client, id_or_name=(name or key_id)) datastore_client.put(entity) entities_to_delete.append(entity) @@ -65,11 +66,14 @@ def test_client_put(datastore_client, entities_to_delete, name, key_id): assert entity.key.id == key_id retrieved_entity = datastore_client.get(entity.key) - # Check the given and retrieved are the the same. + # Check the given and retrieved are the same. assert retrieved_entity == entity -def test_client_put_w_multiple_in_txn(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_put_w_multiple_in_txn( + datastore_client, entities_to_delete, database_id +): with datastore_client.transaction() as xact: entity1 = _get_post(datastore_client) xact.put(entity1) @@ -98,14 +102,18 @@ def test_client_put_w_multiple_in_txn(datastore_client, entities_to_delete): assert len(matches) == 2 -def test_client_query_w_empty_kind(datastore_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_query_w_empty_kind(datastore_client, database_id): query = datastore_client.query(kind="Post") query.ancestor = parent_key(datastore_client) posts = query.fetch(limit=2) assert list(posts) == [] -def test_client_put_w_all_value_types(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_put_w_all_value_types( + datastore_client, entities_to_delete, database_id +): key = datastore_client.key("TestPanObject", 1234) entity = datastore.Entity(key=key) entity["timestamp"] = datetime.datetime(2014, 9, 9, tzinfo=UTC) @@ -127,12 +135,15 @@ def test_client_put_w_all_value_types(datastore_client, entities_to_delete): datastore_client.put(entity) entities_to_delete.append(entity) - # Check the original and retrieved are the the same. + # Check the original and retrieved are the same. retrieved_entity = datastore_client.get(entity.key) assert retrieved_entity == entity -def test_client_put_w_entity_w_self_reference(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_put_w_entity_w_self_reference( + datastore_client, entities_to_delete, database_id +): parent_key = datastore_client.key("Residence", "NewYork") key = datastore_client.key("Person", "name", parent=parent_key) entity = datastore.Entity(key=key) @@ -151,7 +162,8 @@ def test_client_put_w_entity_w_self_reference(datastore_client, entities_to_dele assert stored_persons == [entity] -def test_client_put_w_empty_array(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_client_put_w_empty_array(datastore_client, entities_to_delete, database_id): local_client = _helpers.clone_client(datastore_client) key = local_client.key("EmptyArray", 1234) diff --git a/tests/system/test_query.py b/tests/system/test_query.py index 6b26629f..864bab57 100644 --- a/tests/system/test_query.py +++ b/tests/system/test_query.py @@ -71,7 +71,8 @@ def ancestor_query(query_client, ancestor_key): return _make_ancestor_query(query_client, ancestor_key) -def test_query_w_ancestor(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_ancestor(ancestor_query, database_id): query = ancestor_query expected_matches = 8 @@ -81,7 +82,8 @@ def test_query_w_ancestor(ancestor_query): assert len(entities) == expected_matches -def test_query_w_limit_paging(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_limit_paging(ancestor_query, database_id): query = ancestor_query limit = 5 @@ -101,7 +103,8 @@ def test_query_w_limit_paging(ancestor_query): assert len(new_character_entities) == characters_remaining -def test_query_w_simple_filter(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_simple_filter(ancestor_query, database_id): query = ancestor_query query.add_filter(filter=PropertyFilter("appearances", ">=", 20)) expected_matches = 6 @@ -112,7 +115,8 @@ def test_query_w_simple_filter(ancestor_query): assert len(entities) == expected_matches -def test_query_w_multiple_filters(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_multiple_filters(ancestor_query, database_id): query = ancestor_query query.add_filter(filter=PropertyFilter("appearances", ">=", 26)) query = query.add_filter(filter=PropertyFilter("family", "=", "Stark")) @@ -124,7 +128,8 @@ def test_query_w_multiple_filters(ancestor_query): assert len(entities) == expected_matches -def test_query_key_filter(query_client, ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_key_filter(query_client, ancestor_query, database_id): # Use the client for this test instead of the global. query = ancestor_query rickard_key = query_client.key(*populate_datastore.RICKARD) @@ -137,7 +142,8 @@ def test_query_key_filter(query_client, ancestor_query): assert len(entities) == expected_matches -def test_query_w_order(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_order(ancestor_query, database_id): query = ancestor_query query.order = "appearances" expected_matches = 8 @@ -152,7 +158,8 @@ def test_query_w_order(ancestor_query): assert entities[7]["name"] == populate_datastore.CHARACTERS[3]["name"] -def test_query_w_projection(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_projection(ancestor_query, database_id): filtered_query = ancestor_query filtered_query.projection = ["name", "family"] filtered_query.order = ["name", "family"] @@ -181,7 +188,8 @@ def test_query_w_projection(ancestor_query): assert dict(sansa_entity) == {"name": "Sansa", "family": "Stark"} -def test_query_w_paginate_simple_uuid_keys(query_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_paginate_simple_uuid_keys(query_client, database_id): # See issue #4264 page_query = query_client.query(kind="uuid_key") @@ -199,7 +207,8 @@ def test_query_w_paginate_simple_uuid_keys(query_client): assert page_count > 1 -def test_query_paginate_simple_timestamp_keys(query_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_paginate_simple_timestamp_keys(query_client, database_id): # See issue #4264 page_query = query_client.query(kind="timestamp_key") @@ -217,7 +226,8 @@ def test_query_paginate_simple_timestamp_keys(query_client): assert page_count > 1 -def test_query_w_offset_w_timestamp_keys(query_client): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_offset_w_timestamp_keys(query_client, database_id): # See issue #4675 max_all = 10000 offset = 1 @@ -231,7 +241,8 @@ def test_query_w_offset_w_timestamp_keys(query_client): assert offset_w_limit == all_w_limit[offset:] -def test_query_paginate_with_offset(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_paginate_with_offset(ancestor_query, database_id): page_query = ancestor_query page_query.order = "appearances" offset = 2 @@ -259,7 +270,8 @@ def test_query_paginate_with_offset(ancestor_query): assert entities[2]["name"] == "Arya" -def test_query_paginate_with_start_cursor(query_client, ancestor_key): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_paginate_with_start_cursor(query_client, ancestor_key, database_id): # Don't use fixture, because we need to create a clean copy later. page_query = _make_ancestor_query(query_client, ancestor_key) page_query.order = "appearances" @@ -287,7 +299,8 @@ def test_query_paginate_with_start_cursor(query_client, ancestor_key): assert new_entities[2]["name"] == "Arya" -def test_query_distinct_on(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_distinct_on(ancestor_query, database_id): query = ancestor_query query.distinct_on = ["alive"] expected_matches = 2 @@ -348,7 +361,8 @@ def large_query(large_query_client): (200, populate_datastore.LARGE_CHARACTER_TOTAL_OBJECTS + 1000, 0), ], ) -def test_large_query(large_query, limit, offset, expected): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_large_query(large_query, limit, offset, expected, database_id): page_query = large_query page_query.add_filter(filter=PropertyFilter("family", "=", "Stark")) page_query.add_filter(filter=PropertyFilter("alive", "=", False)) @@ -359,7 +373,8 @@ def test_large_query(large_query, limit, offset, expected): assert len(entities) == expected -def test_query_add_property_filter(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_add_property_filter(ancestor_query, database_id): query = ancestor_query query.add_filter(filter=PropertyFilter("appearances", ">=", 26)) @@ -372,7 +387,8 @@ def test_query_add_property_filter(ancestor_query): assert e["appearances"] >= 26 -def test_query_and_composite_filter(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_and_composite_filter(ancestor_query, database_id): query = ancestor_query query.add_filter( @@ -392,7 +408,8 @@ def test_query_and_composite_filter(ancestor_query): assert entities[0]["name"] == "Jon Snow" -def test_query_or_composite_filter(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_or_composite_filter(ancestor_query, database_id): query = ancestor_query # name = Arya or name = Jon Snow @@ -414,7 +431,8 @@ def test_query_or_composite_filter(ancestor_query): assert entities[1]["name"] == "Jon Snow" -def test_query_add_filters(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_add_filters(ancestor_query, database_id): query = ancestor_query # family = Stark AND name = Jon Snow @@ -430,7 +448,8 @@ def test_query_add_filters(ancestor_query): assert entities[0]["name"] == "Jon Snow" -def test_query_add_complex_filters(ancestor_query): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_add_complex_filters(ancestor_query, database_id): query = ancestor_query # (alive = True OR appearances >= 26) AND (family = Stark) diff --git a/tests/system/test_read_consistency.py b/tests/system/test_read_consistency.py index 9435c5f7..33004352 100644 --- a/tests/system/test_read_consistency.py +++ b/tests/system/test_read_consistency.py @@ -11,12 +11,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +import pytest import time from datetime import datetime, timezone from google.cloud import datastore +from . import _helpers def _parent_key(datastore_client): @@ -33,9 +34,9 @@ def _put_entity(datastore_client, entity_id): return entity -def test_get_w_read_time(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_get_w_read_time(datastore_client, entities_to_delete, database_id): entity = _put_entity(datastore_client, 1) - entities_to_delete.append(entity) # Add some sleep to accommodate server & client clock discrepancy. @@ -62,7 +63,8 @@ def test_get_w_read_time(datastore_client, entities_to_delete): assert retrieved_entity_from_xact["field"] == "old_value" -def test_query_w_read_time(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_query_w_read_time(datastore_client, entities_to_delete, database_id): entity0 = _put_entity(datastore_client, 1) entity1 = _put_entity(datastore_client, 2) entity2 = _put_entity(datastore_client, 3) diff --git a/tests/system/test_transaction.py b/tests/system/test_transaction.py index b380561f..a93538fb 100644 --- a/tests/system/test_transaction.py +++ b/tests/system/test_transaction.py @@ -20,7 +20,10 @@ from . import _helpers -def test_transaction_via_with_statement(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_transaction_via_with_statement( + datastore_client, entities_to_delete, database_id +): key = datastore_client.key("Company", "Google") entity = datastore.Entity(key=key) entity["url"] = "www.google.com" @@ -38,9 +41,9 @@ def test_transaction_via_with_statement(datastore_client, entities_to_delete): assert retrieved_entity == entity +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) def test_transaction_via_explicit_begin_get_commit( - datastore_client, - entities_to_delete, + datastore_client, entities_to_delete, database_id ): # See # github.com/GoogleCloudPlatform/google-cloud-python/issues/1859 @@ -80,7 +83,8 @@ def test_transaction_via_explicit_begin_get_commit( assert after2["balance"] == before_2 + transfer_amount -def test_failure_with_contention(datastore_client, entities_to_delete): +@pytest.mark.parametrize("database_id", [None, _helpers.TEST_DATABASE], indirect=True) +def test_failure_with_contention(datastore_client, entities_to_delete, database_id): contention_prop_name = "baz" local_client = _helpers.clone_client(datastore_client) diff --git a/tests/system/utils/populate_datastore.py b/tests/system/utils/populate_datastore.py index 47394c06..9077241f 100644 --- a/tests/system/utils/populate_datastore.py +++ b/tests/system/utils/populate_datastore.py @@ -123,7 +123,7 @@ def put_objects(count): def add_characters(client=None): if client is None: # Get a client that uses the test dataset. - client = datastore.Client() + client = datastore.Client(database_id="mw-other-db") with client.transaction() as xact: for key_path, character in zip(KEY_PATHS, CHARACTERS): if key_path[-1] != character["name"]: @@ -139,7 +139,7 @@ def add_characters(client=None): def add_uid_keys(client=None): if client is None: # Get a client that uses the test dataset. - client = datastore.Client() + client = datastore.Client(database_id="mw-other-db") num_batches = 2 batch_size = 500 diff --git a/tests/unit/test__http.py b/tests/unit/test__http.py index f9e0a29f..3f292865 100644 --- a/tests/unit/test__http.py +++ b/tests/unit/test__http.py @@ -18,6 +18,8 @@ import pytest import requests +from google.cloud.datastore.helpers import _set_database_id_to_request + def test__make_retry_timeout_kwargs_w_empty(): from google.cloud.datastore._http import _make_retry_timeout_kwargs @@ -97,9 +99,9 @@ def test__make_request_pb_w_instance(): assert foo is passed -def _request_helper(retry=None, timeout=None): +def _request_helper(retry=None, timeout=None, database=None): from google.cloud import _http as connection_module - from google.cloud.datastore._http import _request + from google.cloud.datastore._http import _request, _update_headers project = "PROJECT" method = "METHOD" @@ -113,7 +115,9 @@ def _request_helper(retry=None, timeout=None): kwargs = _retry_timeout_kw(retry, timeout, http) - response = _request(http, project, method, data, base_url, client_info, **kwargs) + response = _request( + http, project, method, data, base_url, client_info, database=database, **kwargs + ) assert response == response_data # Check that the mocks were called as expected. @@ -123,7 +127,7 @@ def _request_helper(retry=None, timeout=None): "User-Agent": user_agent, connection_module.CLIENT_INFO_HEADER: user_agent, } - + _update_headers(expected_headers, project, database_id=database) if retry is not None: retry.assert_called_once_with(http.request) @@ -133,18 +137,21 @@ def _request_helper(retry=None, timeout=None): ) -def test__request_defaults(): - _request_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test__request_defaults(database_id): + _request_helper(database=database_id) -def test__request_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test__request_w_retry(database_id): retry = mock.MagicMock() - _request_helper(retry=retry) + _request_helper(retry=retry, database=database_id) -def test__request_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test__request_w_timeout(database_id): timeout = 5.0 - _request_helper(timeout=timeout) + _request_helper(timeout=timeout, database=database_id) def test__request_failure(): @@ -175,7 +182,7 @@ def test__request_failure(): assert exc.match(expected_message) -def _rpc_helper(retry=None, timeout=None): +def _rpc_helper(retry=None, timeout=None, database=None): from google.cloud.datastore._http import _rpc from google.cloud.datastore_v1.types import datastore as datastore_pb2 @@ -203,7 +210,8 @@ def _rpc_helper(retry=None, timeout=None): client_info, request_pb, datastore_pb2.BeginTransactionResponse, - **kwargs + database=database, + **kwargs, ) assert result == response_pb._pb @@ -215,22 +223,26 @@ def _rpc_helper(retry=None, timeout=None): request_pb._pb.SerializeToString(), base_url, client_info, - **kwargs + database=database, + **kwargs, ) -def test__rpc_defaults(): - _rpc_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test__rpc_defaults(database_id): + _rpc_helper(database=database_id) -def test__rpc_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test__rpc_w_retry(database_id): retry = mock.MagicMock() - _rpc_helper(retry=retry) + _rpc_helper(retry=retry, database=database_id) -def test__rpc_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test__rpc_w_timeout(database_id): timeout = 5.0 - _rpc_helper(timeout=timeout) + _rpc_helper(timeout=timeout, database=database_id) def test_api_ctor(): @@ -245,6 +257,7 @@ def _lookup_single_helper( empty=True, retry=None, timeout=None, + database=None, ): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 @@ -283,6 +296,7 @@ def _lookup_single_helper( "keys": [key_pb], "read_options": read_options, } + _set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.lookup(request=request, **kwargs) @@ -304,6 +318,8 @@ def _lookup_single_helper( datastore_pb2.LookupRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) if retry is not None: @@ -344,11 +360,7 @@ def test_api_lookup_single_key_hit_w_timeout(): def _lookup_multiple_helper( - found=0, - missing=0, - deferred=0, - retry=None, - timeout=None, + found=0, missing=0, deferred=0, retry=None, timeout=None, database=None ): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 @@ -416,6 +428,8 @@ def _lookup_multiple_helper( datastore_pb2.LookupRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert list(request.keys) == [key_pb1._pb, key_pb2._pb] assert request.read_options == read_options._pb @@ -454,6 +468,7 @@ def _run_query_helper( found=0, retry=None, timeout=None, + database=None, ): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 @@ -520,6 +535,8 @@ def _run_query_helper( datastore_pb2.RunQueryRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert request.partition_id == partition_id._pb assert request.query == query_pb._pb @@ -558,9 +575,7 @@ def test_api_run_query_w_namespace_nonempty_result(): def _run_aggregation_query_helper( - transaction=None, - retry=None, - timeout=None, + transaction=None, retry=None, timeout=None, database=None ): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 @@ -623,6 +638,8 @@ def _run_aggregation_query_helper( datastore_pb2.RunAggregationQueryRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert request.partition_id == partition_id._pb @@ -649,7 +666,7 @@ def test_api_run_aggregation_query_w_transaction(): _run_aggregation_query_helper(transaction=transaction) -def _begin_transaction_helper(options=None, retry=None, timeout=None): +def _begin_transaction_helper(options=None, retry=None, timeout=None, database=None): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" @@ -672,7 +689,7 @@ def _begin_transaction_helper(options=None, retry=None, timeout=None): # Make request. ds_api = _make_http_datastore_api(client) request = {"project_id": project} - + _set_database_id_to_request(request, database) if options is not None: request["transaction_options"] = options @@ -690,37 +707,43 @@ def _begin_transaction_helper(options=None, retry=None, timeout=None): datastore_pb2.BeginTransactionRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) -def test_api_begin_transaction_wo_options(): - _begin_transaction_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_begin_transaction_wo_options(database_id): + _begin_transaction_helper(database=database_id) -def test_api_begin_transaction_w_options(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_begin_transaction_w_options(database_id): from google.cloud.datastore_v1.types import TransactionOptions read_only = TransactionOptions.ReadOnly._meta.pb() options = TransactionOptions(read_only=read_only) - _begin_transaction_helper(options=options) + _begin_transaction_helper(options=options, database=database_id) -def test_api_begin_transaction_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_begin_transaction_w_retry(database_id): retry = mock.MagicMock() - _begin_transaction_helper(retry=retry) + _begin_transaction_helper(retry=retry, database=database_id) -def test_api_begin_transaction_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_begin_transaction_w_timeout(database_id): timeout = 5.0 - _begin_transaction_helper(timeout=timeout) + _begin_transaction_helper(timeout=timeout, database=database_id) -def _commit_helper(transaction=None, retry=None, timeout=None): +def _commit_helper(transaction=None, retry=None, timeout=None, database=None): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore.helpers import _new_value_pb project = "PROJECT" - key_pb = _make_key_pb(project) + key_pb = _make_key_pb(project, database=database) rsp_pb = datastore_pb2.CommitResponse() req_pb = datastore_pb2.CommitRequest() mutation = req_pb._pb.mutations.add() @@ -744,7 +767,7 @@ def _commit_helper(transaction=None, retry=None, timeout=None): ds_api = _make_http_datastore_api(client) request = {"project_id": project, "mutations": [mutation]} - + _set_database_id_to_request(request, database) if transaction is not None: request["transaction"] = transaction mode = request["mode"] = rq_class.Mode.TRANSACTIONAL @@ -764,6 +787,8 @@ def _commit_helper(transaction=None, retry=None, timeout=None): rq_class(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert list(request.mutations) == [mutation] assert request.mode == mode @@ -774,27 +799,31 @@ def _commit_helper(transaction=None, retry=None, timeout=None): assert request.transaction == b"" -def test_api_commit_wo_transaction(): - _commit_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_commit_wo_transaction(database_id): + _commit_helper(database=database_id) -def test_api_commit_w_transaction(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_commit_w_transaction(database_id): transaction = b"xact" - _commit_helper(transaction=transaction) + _commit_helper(transaction=transaction, database=database_id) -def test_api_commit_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_commit_w_retry(database_id): retry = mock.MagicMock() - _commit_helper(retry=retry) + _commit_helper(retry=retry, database=database_id) -def test_api_commit_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_commit_w_timeout(database_id): timeout = 5.0 - _commit_helper(timeout=timeout) + _commit_helper(timeout=timeout, database=database_id) -def _rollback_helper(retry=None, timeout=None): +def _rollback_helper(retry=None, timeout=None, database=None): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" @@ -816,6 +845,7 @@ def _rollback_helper(retry=None, timeout=None): # Make request. ds_api = _make_http_datastore_api(client) request = {"project_id": project, "transaction": transaction} + _set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.rollback(request=request, **kwargs) @@ -830,25 +860,30 @@ def _rollback_helper(retry=None, timeout=None): datastore_pb2.RollbackRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert request.transaction == transaction -def test_api_rollback_ok(): - _rollback_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_rollback_ok(database_id): + _rollback_helper(database=database_id) -def test_api_rollback_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_rollback_w_retry(database_id): retry = mock.MagicMock() - _rollback_helper(retry=retry) + _rollback_helper(retry=retry, database=database_id) -def test_api_rollback_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_rollback_w_timeout(database_id): timeout = 5.0 - _rollback_helper(timeout=timeout) + _rollback_helper(timeout=timeout, database=database_id) -def _allocate_ids_helper(count=0, retry=None, timeout=None): +def _allocate_ids_helper(count=0, retry=None, timeout=None, database=None): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" @@ -857,9 +892,9 @@ def _allocate_ids_helper(count=0, retry=None, timeout=None): rsp_pb = datastore_pb2.AllocateIdsResponse() for i_count in range(count): - requested = _make_key_pb(project, id_=None) + requested = _make_key_pb(project, id_=None, database=database) before_key_pbs.append(requested) - allocated = _make_key_pb(project, id_=i_count) + allocated = _make_key_pb(project, id_=i_count, database=database) after_key_pbs.append(allocated) rsp_pb._pb.keys.add().CopyFrom(allocated._pb) @@ -876,6 +911,7 @@ def _allocate_ids_helper(count=0, retry=None, timeout=None): ds_api = _make_http_datastore_api(client) request = {"project_id": project, "keys": before_key_pbs} + _set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.allocate_ids(request=request, **kwargs) @@ -890,31 +926,37 @@ def _allocate_ids_helper(count=0, retry=None, timeout=None): datastore_pb2.AllocateIdsRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert len(request.keys) == len(before_key_pbs) for key_before, key_after in zip(before_key_pbs, request.keys): assert key_before == key_after -def test_api_allocate_ids_empty(): - _allocate_ids_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_allocate_ids_empty(database_id): + _allocate_ids_helper(database=database_id) -def test_api_allocate_ids_non_empty(): - _allocate_ids_helper(count=2) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_allocate_ids_non_empty(database_id): + _allocate_ids_helper(count=2, database=database_id) -def test_api_allocate_ids_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_allocate_ids_w_retry(database_id): retry = mock.MagicMock() - _allocate_ids_helper(retry=retry) + _allocate_ids_helper(retry=retry, database=database_id) -def test_api_allocate_ids_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_allocate_ids_w_timeout(database_id): timeout = 5.0 - _allocate_ids_helper(timeout=timeout) + _allocate_ids_helper(timeout=timeout, database=database_id) -def _reserve_ids_helper(count=0, retry=None, timeout=None): +def _reserve_ids_helper(count=0, retry=None, timeout=None, database=None): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" @@ -922,7 +964,7 @@ def _reserve_ids_helper(count=0, retry=None, timeout=None): rsp_pb = datastore_pb2.ReserveIdsResponse() for i_count in range(count): - requested = _make_key_pb(project, id_=i_count) + requested = _make_key_pb(project, id_=i_count, database=database) before_key_pbs.append(requested) http = _make_requests_session( @@ -938,6 +980,7 @@ def _reserve_ids_helper(count=0, retry=None, timeout=None): ds_api = _make_http_datastore_api(client) request = {"project_id": project, "keys": before_key_pbs} + _set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.reserve_ids(request=request, **kwargs) @@ -951,28 +994,56 @@ def _reserve_ids_helper(count=0, retry=None, timeout=None): datastore_pb2.AllocateIdsRequest(), retry=retry, timeout=timeout, + project=project, + database=database, ) assert len(request.keys) == len(before_key_pbs) for key_before, key_after in zip(before_key_pbs, request.keys): assert key_before == key_after -def test_api_reserve_ids_empty(): - _reserve_ids_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_reserve_ids_empty(database_id): + _reserve_ids_helper(database=database_id) -def test_api_reserve_ids_non_empty(): - _reserve_ids_helper(count=2) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_reserve_ids_non_empty(database_id): + _reserve_ids_helper(count=2, database=database_id) -def test_api_reserve_ids_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_reserve_ids_w_retry(database_id): retry = mock.MagicMock() - _reserve_ids_helper(retry=retry) + _reserve_ids_helper(retry=retry, database=database_id) -def test_api_reserve_ids_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_api_reserve_ids_w_timeout(database_id): timeout = 5.0 - _reserve_ids_helper(timeout=timeout) + _reserve_ids_helper(timeout=timeout, database=database_id) + + +def test_update_headers_without_database_id(): + from google.cloud.datastore._http import _update_headers + + headers = {} + project_id = "someproject" + _update_headers(headers, project_id) + assert headers["x-goog-request-params"] == f"project_id={project_id}" + + +def test_update_headers_with_database_id(): + from google.cloud.datastore._http import _update_headers + + headers = {} + project_id = "someproject" + database_id = "somedb" + _update_headers(headers, project_id, database_id=database_id) + assert ( + headers["x-goog-request-params"] + == f"project_id={project_id}&database_id={database_id}" + ) def _make_http_datastore_api(*args, **kwargs): @@ -1002,13 +1073,13 @@ def _build_expected_url(api_base_url, project, method): return "/".join([api_base_url, API_VERSION, "projects", project + ":" + method]) -def _make_key_pb(project, id_=1234): +def _make_key_pb(project, id_=1234, database=None): from google.cloud.datastore.key import Key path_args = ("Kind",) if id_ is not None: path_args += (id_,) - return Key(*path_args, project=project).to_protobuf() + return Key(*path_args, project=project, database=database).to_protobuf() _USER_AGENT = "TESTING USER AGENT" @@ -1022,15 +1093,18 @@ def _make_client_info(user_agent=_USER_AGENT): return client_info -def _verify_protobuf_call(http, expected_url, pb, retry=None, timeout=None): +def _verify_protobuf_call( + http, expected_url, pb, retry=None, timeout=None, project=None, database=None +): from google.cloud import _http as connection_module + from google.cloud.datastore._http import _update_headers expected_headers = { "Content-Type": "application/x-protobuf", "User-Agent": _USER_AGENT, connection_module.CLIENT_INFO_HEADER: _USER_AGENT, } - + _update_headers(expected_headers, project, database_id=database) if retry is not None: retry.assert_called_once_with(http.request) diff --git a/tests/unit/test_aggregation.py b/tests/unit/test_aggregation.py index 59c07b27..a45b05b2 100644 --- a/tests/unit/test_aggregation.py +++ b/tests/unit/test_aggregation.py @@ -16,6 +16,7 @@ import pytest from google.cloud.datastore.aggregation import CountAggregation, AggregationQuery +from google.cloud.datastore.helpers import _set_database_id_to_request from tests.unit.test_query import _make_query, _make_client @@ -34,11 +35,17 @@ def test_count_aggregation_to_pb(): @pytest.fixture -def client(): - return _make_client() +def database_id(request): + return request.param -def test_project(client): +@pytest.fixture +def client(database_id): + return _make_client(database=database_id) + + +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_project(client, database_id): # Fallback to client query = _make_query(client) aggregation_query = _make_aggregation_query(client=client, query=query) @@ -50,21 +57,8 @@ def test_project(client): assert aggregation_query.project == "other-project" -def test_database(client): - # Fallback to client - client.database = None - query = _make_query(client, database=None) - client.database = "other-database" - aggregation_query = _make_aggregation_query(client=client, query=query) - assert aggregation_query.database == "other-database" - - # Fallback to query - query = _make_query(client, database="third-database") - aggregation_query = _make_aggregation_query(client=client, query=query) - assert aggregation_query.database == "third-database" - - -def test_namespace(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_namespace(client, database_id): # Fallback to client client.namespace = "other-namespace" query = _make_query(client) @@ -77,7 +71,8 @@ def test_namespace(client): assert aggregation_query.namespace == "third-namespace" -def test_pb_over_query(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_pb_over_query(client, database_id): from google.cloud.datastore.query import _pb_from_query query = _make_query(client) @@ -87,7 +82,8 @@ def test_pb_over_query(client): assert pb.aggregations == [] -def test_pb_over_query_with_count(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_pb_over_query_with_count(client, database_id): from google.cloud.datastore.query import _pb_from_query query = _make_query(client) @@ -100,7 +96,8 @@ def test_pb_over_query_with_count(client): assert pb.aggregations[0] == CountAggregation(alias="total")._to_pb() -def test_pb_over_query_with_add_aggregation(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_pb_over_query_with_add_aggregation(client, database_id): from google.cloud.datastore.query import _pb_from_query query = _make_query(client) @@ -113,7 +110,8 @@ def test_pb_over_query_with_add_aggregation(client): assert pb.aggregations[0] == CountAggregation(alias="total")._to_pb() -def test_pb_over_query_with_add_aggregations(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_pb_over_query_with_add_aggregations(client, database_id): from google.cloud.datastore.query import _pb_from_query aggregations = [ @@ -132,7 +130,8 @@ def test_pb_over_query_with_add_aggregations(client): assert pb.aggregations[1] == CountAggregation(alias="all")._to_pb() -def test_query_fetch_defaults_w_client_attr(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_query_fetch_defaults_w_client_attr(client, database_id): from google.cloud.datastore.aggregation import AggregationResultIterator query = _make_query(client) @@ -146,10 +145,11 @@ def test_query_fetch_defaults_w_client_attr(client): assert iterator._timeout is None -def test_query_fetch_w_explicit_client_w_retry_w_timeout(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_query_fetch_w_explicit_client_w_retry_w_timeout(client, database_id): from google.cloud.datastore.aggregation import AggregationResultIterator - other_client = _make_client() + other_client = _make_client(database=database_id) query = _make_query(client) aggregation_query = _make_aggregation_query(client=client, query=query) retry = mock.Mock() @@ -166,10 +166,11 @@ def test_query_fetch_w_explicit_client_w_retry_w_timeout(client): assert iterator._timeout == timeout -def test_query_fetch_w_explicit_client_w_limit(client): +@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True) +def test_query_fetch_w_explicit_client_w_limit(client, database_id): from google.cloud.datastore.aggregation import AggregationResultIterator - other_client = _make_client() + other_client = _make_client(database=database_id) query = _make_query(client) aggregation_query = _make_aggregation_query(client=client, query=query) limit = 2 @@ -339,7 +340,7 @@ def test_iterator__next_page_no_more(): ds_api.run_aggregation_query.assert_not_called() -def _next_page_helper(txn_id=None, retry=None, timeout=None): +def _next_page_helper(txn_id=None, retry=None, timeout=None, database_id=None): from google.api_core import page_iterator from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 @@ -357,10 +358,12 @@ def _next_page_helper(txn_id=None, retry=None, timeout=None): project = "prujekt" ds_api = _make_datastore_api_for_aggregation(result_1, result_2) if txn_id is None: - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) else: transaction = mock.Mock(id=txn_id, spec=["id"]) - client = _Client(project, datastore_api=ds_api, transaction=transaction) + client = _Client( + project, datastore_api=ds_api, transaction=transaction, database=database_id + ) query = _make_query(client) kwargs = {} @@ -389,14 +392,15 @@ def _next_page_helper(txn_id=None, retry=None, timeout=None): aggregation_query = AggregationQuery(client=client, query=query) assert ds_api.run_aggregation_query.call_count == 2 + expected_request = { + "project_id": project, + "partition_id": partition_id, + "read_options": read_options, + "aggregation_query": aggregation_query._to_pb(), + } + _set_database_id_to_request(expected_request, database_id) expected_call = mock.call( - request={ - "project_id": project, - "database_id": "", - "partition_id": partition_id, - "read_options": read_options, - "aggregation_query": aggregation_query._to_pb(), - }, + request=expected_request, **kwargs, ) assert ds_api.run_aggregation_query.call_args_list == ( @@ -430,7 +434,7 @@ def __init__( namespace=None, transaction=None, *, - database="", + database=None, ): self.project = project self.database = database diff --git a/tests/unit/test_batch.py b/tests/unit/test_batch.py index f3fc049e..25a968fc 100644 --- a/tests/unit/test_batch.py +++ b/tests/unit/test_batch.py @@ -18,6 +18,8 @@ import mock import pytest +from google.cloud.datastore.helpers import _set_database_id_to_request + def _make_batch(client): from google.cloud.datastore.batch import Batch @@ -25,16 +27,16 @@ def _make_batch(client): return Batch(client) -def test_batch_ctor(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_ctor(database_id): project = "PROJECT" - database = "DATABASE" namespace = "NAMESPACE" - client = _Client(project, database=database, namespace=namespace) + client = _Client(project, database=database_id, namespace=namespace) batch = _make_batch(client) assert batch.project == project assert batch._client is client - assert batch.database == database + assert batch.database == database_id assert batch.namespace == namespace assert batch._id is None assert batch._status == batch._INITIAL @@ -42,12 +44,12 @@ def test_batch_ctor(): assert batch._partial_key_entities == [] -def test_batch_current(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_current(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" - database = "DATABASE" - client = _Client(project, database=database) + client = _Client(project, database=database_id) batch1 = _make_batch(client) batch2 = _make_batch(client) @@ -71,20 +73,20 @@ def test_batch_current(): commit_method = client._datastore_api.commit assert commit_method.call_count == 2 mode = datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL - commit_method.assert_called_with( - request={ - "project_id": project, - "database_id": database, - "mode": mode, - "mutations": [], - "transaction": None, - } - ) - - -def test_batch_put_w_entity_wo_key(): + expected_request = { + "project_id": project, + "mode": mode, + "mutations": [], + "transaction": None, + } + _set_database_id_to_request(expected_request, database_id) + commit_method.assert_called_with(request=expected_request) + + +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_put_w_entity_wo_key(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) entity = _Entity() @@ -93,24 +95,26 @@ def test_batch_put_w_entity_wo_key(): batch.put(entity) -def test_batch_put_w_wrong_status(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_put_w_wrong_status(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) entity = _Entity() - entity.key = _Key(project=project) + entity.key = _Key(project=project, database=database_id) assert batch._status == batch._INITIAL with pytest.raises(ValueError): batch.put(entity) -def test_batch_put_w_key_wrong_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_put_w_key_wrong_project(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) entity = _Entity() - entity.key = _Key(project="OTHER") + entity.key = _Key(project="OTHER", database=database_id) batch.begin() with pytest.raises(ValueError): @@ -119,24 +123,24 @@ def test_batch_put_w_key_wrong_project(): def test_batch_put_w_key_wrong_database(): project = "PROJECT" - database = "DATABASE" - client = _Client(project, database=database) + client = _Client(project) batch = _make_batch(client) entity = _Entity() - entity.key = _Key(project=project, database=None) + entity.key = _Key(project=project, database="somedb") batch.begin() with pytest.raises(ValueError): batch.put(entity) -def test_batch_put_w_entity_w_partial_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_put_w_entity_w_partial_key(database_id): project = "PROJECT" properties = {"foo": "bar"} - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) entity = _Entity(properties) - key = entity.key = _Key(project) + key = entity.key = _Key(project, database=database_id) key._id = None batch.begin() @@ -147,14 +151,15 @@ def test_batch_put_w_entity_w_partial_key(): assert batch._partial_key_entities == [entity] -def test_batch_put_w_entity_w_completed_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_put_w_entity_w_completed_key(database_id): project = "PROJECT" properties = {"foo": "bar", "baz": "qux", "spam": [1, 2, 3], "frotz": []} - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) entity = _Entity(properties) entity.exclude_from_indexes = ("baz", "spam") - key = entity.key = _Key(project) + key = entity.key = _Key(project, database=database_id) batch.begin() batch.put(entity) @@ -175,11 +180,12 @@ def test_batch_put_w_entity_w_completed_key(): assert "frotz" in prop_dict -def test_batch_delete_w_wrong_status(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_delete_w_wrong_status(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) - key = _Key(project=project) + key = _Key(project=project, database=database_id) key._id = None assert batch._status == batch._INITIAL @@ -188,11 +194,12 @@ def test_batch_delete_w_wrong_status(): batch.delete(key) -def test_batch_delete_w_partial_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_delete_w_partial_key(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) - key = _Key(project=project) + key = _Key(project=project, database=database_id) key._id = None batch.begin() @@ -201,11 +208,12 @@ def test_batch_delete_w_partial_key(): batch.delete(key) -def test_batch_delete_w_key_wrong_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_delete_w_key_wrong_project(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) - key = _Key(project="OTHER") + key = _Key(project="OTHER", database=database_id) batch.begin() with pytest.raises(ValueError): @@ -224,11 +232,12 @@ def test_batch_delete_w_key_wrong_database(): batch.delete(key) -def test_batch_delete_w_completed_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_delete_w_completed_key(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) - key = _Key(project) + key = _Key(project, database=database_id) batch.begin() batch.delete(key) @@ -237,9 +246,10 @@ def test_batch_delete_w_completed_key(): assert mutated_key == key._key -def test_batch_begin_w_wrong_status(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_begin_w_wrong_status(database_id): project = "PROJECT" - client = _Client(project, None) + client = _Client(project, database=database_id) batch = _make_batch(client) batch._status = batch._IN_PROGRESS @@ -247,9 +257,10 @@ def test_batch_begin_w_wrong_status(): batch.begin() -def test_batch_begin(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_begin(database_id): project = "PROJECT" - client = _Client(project, None) + client = _Client(project, database=database_id) batch = _make_batch(client) assert batch._status == batch._INITIAL @@ -258,9 +269,10 @@ def test_batch_begin(): assert batch._status == batch._IN_PROGRESS -def test_batch_rollback_w_wrong_status(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_rollback_w_wrong_status(database_id): project = "PROJECT" - client = _Client(project, None) + client = _Client(project, database=database_id) batch = _make_batch(client) assert batch._status == batch._INITIAL @@ -268,9 +280,10 @@ def test_batch_rollback_w_wrong_status(): batch.rollback() -def test_batch_rollback(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_rollback(database_id): project = "PROJECT" - client = _Client(project, None) + client = _Client(project, database=database_id) batch = _make_batch(client) batch.begin() assert batch._status == batch._IN_PROGRESS @@ -280,9 +293,10 @@ def test_batch_rollback(): assert batch._status == batch._ABORTED -def test_batch_commit_wrong_status(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_commit_wrong_status(database_id): project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database_id) batch = _make_batch(client) assert batch._status == batch._INITIAL @@ -290,11 +304,11 @@ def test_batch_commit_wrong_status(): batch.commit() -def _batch_commit_helper(timeout=None, retry=None): +def _batch_commit_helper(timeout=None, retry=None, database=None): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" - client = _Client(project) + client = _Client(project, database=database) batch = _make_batch(client) assert batch._status == batch._INITIAL @@ -314,39 +328,41 @@ def _batch_commit_helper(timeout=None, retry=None): commit_method = client._datastore_api.commit mode = datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL - commit_method.assert_called_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": [], - "transaction": None, - }, - **kwargs - ) + expected_request = { + "project_id": project, + "mode": mode, + "mutations": [], + "transaction": None, + } + _set_database_id_to_request(expected_request, database) + commit_method.assert_called_with(request=expected_request, **kwargs) -def test_batch_commit(): - _batch_commit_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_commit(database_id): + _batch_commit_helper(database=database_id) -def test_batch_commit_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_commit_w_timeout(database_id): timeout = 100000 - _batch_commit_helper(timeout=timeout) + _batch_commit_helper(timeout=timeout, database=database_id) -def test_batch_commit_w_retry(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_commit_w_retry(database_id): retry = mock.Mock(spec=[]) - _batch_commit_helper(retry=retry) + _batch_commit_helper(retry=retry, database=database_id) -def test_batch_commit_w_partial_key_entity(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_commit_w_partial_key_entity(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" new_id = 1234 ds_api = _make_datastore_api(new_id) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) batch = _make_batch(client) entity = _Entity({}) key = entity.key = _Key(project) @@ -361,28 +377,29 @@ def test_batch_commit_w_partial_key_entity(): assert batch._status == batch._FINISHED mode = datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL - ds_api.commit.assert_called_once_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": [], - "transaction": None, - } - ) + expected_request = { + "project_id": project, + "mode": mode, + "mutations": [], + "transaction": None, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.commit.assert_called_once_with(request=expected_request) + assert not entity.key.is_partial assert entity.key._id == new_id -def test_batch_as_context_mgr_wo_error(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_as_context_mgr_wo_error(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" properties = {"foo": "bar"} entity = _Entity(properties) - key = entity.key = _Key(project) + key = entity.key = _Key(project, database=database_id) - client = _Client(project) + client = _Client(project, database=database_id) assert list(client._batches) == [] with _make_batch(client) as batch: @@ -396,28 +413,28 @@ def test_batch_as_context_mgr_wo_error(): commit_method = client._datastore_api.commit mode = datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL - commit_method.assert_called_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": batch.mutations, - "transaction": None, - } - ) - - -def test_batch_as_context_mgr_nested(): + expected_request = { + "project_id": project, + "mode": mode, + "mutations": batch.mutations, + "transaction": None, + } + _set_database_id_to_request(expected_request, database_id) + commit_method.assert_called_with(request=expected_request) + + +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_as_context_mgr_nested(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" properties = {"foo": "bar"} entity1 = _Entity(properties) - key1 = entity1.key = _Key(project) + key1 = entity1.key = _Key(project, database=database_id) entity2 = _Entity(properties) - key2 = entity2.key = _Key(project) + key2 = entity2.key = _Key(project, database=database_id) - client = _Client(project) + client = _Client(project, database=database_id) assert list(client._batches) == [] with _make_batch(client) as batch1: @@ -442,33 +459,33 @@ def test_batch_as_context_mgr_nested(): assert commit_method.call_count == 2 mode = datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL - commit_method.assert_called_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": batch1.mutations, - "transaction": None, - } - ) - commit_method.assert_called_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": batch2.mutations, - "transaction": None, - } - ) - - -def test_batch_as_context_mgr_w_error(): + expected_request_1 = { + "project_id": project, + "mode": mode, + "mutations": batch1.mutations, + "transaction": None, + } + expected_request_2 = { + "project_id": project, + "mode": mode, + "mutations": batch1.mutations, + "transaction": None, + } + _set_database_id_to_request(expected_request_1, database_id) + _set_database_id_to_request(expected_request_2, database_id) + + commit_method.assert_called_with(request=expected_request_1) + commit_method.assert_called_with(request=expected_request_2) + + +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_batch_as_context_mgr_w_error(database_id): project = "PROJECT" properties = {"foo": "bar"} entity = _Entity(properties) - key = entity.key = _Key(project) + key = entity.key = _Key(project, database=database_id) - client = _Client(project) + client = _Client(project, database=database_id) assert list(client._batches) == [] try: @@ -544,7 +561,7 @@ class _Key(object): _id = 1234 _stored = None - def __init__(self, project, database=""): + def __init__(self, project, database=None): self.project = project self.database = database @@ -574,7 +591,7 @@ def completed_key(self, new_id): class _Client(object): - def __init__(self, project, datastore_api=None, namespace=None, database=""): + def __init__(self, project, datastore_api=None, namespace=None, database=None): self.project = project if datastore_api is None: datastore_api = _make_datastore_api() diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 55b7200b..06cfa37f 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -18,6 +18,8 @@ import mock import pytest +from google.cloud.datastore.helpers import _set_database_id_to_request + PROJECT = "dummy-project-123" DATABASE = "dummy-database-123" @@ -126,7 +128,8 @@ def test_client_ctor_w_project_no_environ(): _make_client(project=None) -def test_client_ctor_w_implicit_inputs(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_ctor_w_implicit_inputs(database_id): from google.cloud.datastore.client import Client from google.cloud.datastore.client import _CLIENT_INFO from google.cloud.datastore.client import _DATASTORE_BASE_URL @@ -142,10 +145,10 @@ def test_client_ctor_w_implicit_inputs(): with patch1 as _determine_default_project: with patch2 as default: - client = Client() + client = Client(database=database_id) assert client.project == other - assert client.database == "" + assert client.database == database_id assert client.namespace is None assert client._credentials is creds assert client._client_info is _CLIENT_INFO @@ -162,7 +165,8 @@ def test_client_ctor_w_implicit_inputs(): _determine_default_project.assert_called_once_with(None) -def test_client_ctor_w_explicit_inputs(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_ctor_w_explicit_inputs(database_id): from google.api_core.client_options import ClientOptions other = "other" @@ -192,7 +196,8 @@ def test_client_ctor_w_explicit_inputs(): assert list(client._batch_stack) == [] -def test_client_ctor_use_grpc_default(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_ctor_use_grpc_default(database_id): import google.cloud.datastore.client as MUT project = "PROJECT" @@ -200,20 +205,32 @@ def test_client_ctor_use_grpc_default(): http = object() with mock.patch.object(MUT, "_USE_GRPC", new=True): - client1 = _make_client(project=PROJECT, credentials=creds, _http=http) + client1 = _make_client( + project=PROJECT, credentials=creds, _http=http, database=database_id + ) assert client1._use_grpc # Explicitly over-ride the environment. client2 = _make_client( - project=project, credentials=creds, _http=http, _use_grpc=False + project=project, + credentials=creds, + _http=http, + _use_grpc=False, + database=database_id, ) assert not client2._use_grpc with mock.patch.object(MUT, "_USE_GRPC", new=False): - client3 = _make_client(project=PROJECT, credentials=creds, _http=http) + client3 = _make_client( + project=PROJECT, credentials=creds, _http=http, database=database_id + ) assert not client3._use_grpc # Explicitly over-ride the environment. client4 = _make_client( - project=project, credentials=creds, _http=http, _use_grpc=True + project=project, + credentials=creds, + _http=http, + _use_grpc=True, + database=database_id, ) assert client4._use_grpc @@ -414,12 +431,13 @@ def test_client_get_multi_no_keys(): ds_api.lookup.assert_not_called() -def test_client_get_multi_miss(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_miss(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore.key import Key creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) ds_api = _make_datastore_api() client._datastore_api_internal = ds_api @@ -428,17 +446,17 @@ def test_client_get_multi_miss(): assert results == [] read_options = datastore_pb2.ReadOptions() - ds_api.lookup.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key.to_protobuf()], - "read_options": read_options, - } - ) + expected_request = { + "project_id": PROJECT, + "keys": [key.to_protobuf()], + "read_options": read_options, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.lookup.assert_called_once_with(request=expected_request) -def test_client_get_multi_miss_w_missing(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_miss_w_missing(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 from google.cloud.datastore.key import Key @@ -449,18 +467,19 @@ def test_client_get_multi_miss_w_missing(): # Make a missing entity pb to be returned from mock backend. missed = entity_pb2.Entity() missed.key.partition_id.project_id = PROJECT + missed.key.partition_id.database_id = database_id path_element = missed._pb.key.path.add() path_element.kind = KIND path_element.id = ID creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) # Set missing entity on mock connection. lookup_response = _make_lookup_response(missing=[missed._pb]) ds_api = _make_datastore_api(lookup_response=lookup_response) client._datastore_api_internal = ds_api - key = Key(KIND, ID, project=PROJECT) + key = Key(KIND, ID, project=PROJECT, database=database_id) missing = [] entities = client.get_multi([key], missing=missing) assert entities == [] @@ -468,14 +487,13 @@ def test_client_get_multi_miss_w_missing(): assert [missed.key.to_protobuf() for missed in missing] == [key_pb._pb] read_options = datastore_pb2.ReadOptions() - ds_api.lookup.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key_pb], - "read_options": read_options, - } - ) + expected_request = { + "project_id": PROJECT, + "keys": [key_pb], + "read_options": read_options, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.lookup.assert_called_once_with(request=expected_request) def test_client_get_multi_w_missing_non_empty(): @@ -502,16 +520,17 @@ def test_client_get_multi_w_deferred_non_empty(): client.get_multi([key], deferred=deferred) -def test_client_get_multi_miss_w_deferred(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_miss_w_deferred(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore.key import Key - key = Key("Kind", 1234, project=PROJECT) + key = Key("Kind", 1234, project=PROJECT, database=database_id) key_pb = key.to_protobuf() # Set deferred entity on mock connection. creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) lookup_response = _make_lookup_response(deferred=[key_pb]) ds_api = _make_datastore_api(lookup_response=lookup_response) client._datastore_api_internal = ds_api @@ -520,27 +539,27 @@ def test_client_get_multi_miss_w_deferred(): entities = client.get_multi([key], deferred=deferred) assert entities == [] assert [def_key.to_protobuf() for def_key in deferred] == [key_pb] - read_options = datastore_pb2.ReadOptions() - ds_api.lookup.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key_pb], - "read_options": read_options, - } - ) + expected_request = { + "project_id": PROJECT, + "keys": [key_pb], + "read_options": read_options, + } + _set_database_id_to_request(expected_request, database_id) + + ds_api.lookup.assert_called_once_with(request=expected_request) -def test_client_get_multi_w_deferred_from_backend_but_not_passed(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_w_deferred_from_backend_but_not_passed(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 from google.cloud.datastore.entity import Entity from google.cloud.datastore.key import Key - key1 = Key("Kind", project=PROJECT) + key1 = Key("Kind", project=PROJECT, database=database_id) key1_pb = key1.to_protobuf() - key2 = Key("Kind", 2345, project=PROJECT) + key2 = Key("Kind", 2345, project=PROJECT, database=database_id) key2_pb = key2.to_protobuf() entity1_pb = entity_pb2.Entity() @@ -549,7 +568,7 @@ def test_client_get_multi_w_deferred_from_backend_but_not_passed(): entity2_pb._pb.key.CopyFrom(key2_pb._pb) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) # Mock up two separate requests. Using an iterable as side_effect # allows multiple return values. lookup_response1 = _make_lookup_response(results=[entity1_pb], deferred=[key2_pb]) @@ -567,34 +586,39 @@ def test_client_get_multi_w_deferred_from_backend_but_not_passed(): assert isinstance(found[0], Entity) assert found[0].key.path == key1.path assert found[0].key.project == key1.project + assert found[0].key.database == key1.database assert isinstance(found[1], Entity) assert found[1].key.path == key2.path assert found[1].key.project == key2.project + assert found[1].key.database == key2.database assert ds_api.lookup.call_count == 2 read_options = datastore_pb2.ReadOptions() + expected_request_1 = { + "project_id": PROJECT, + "keys": [key2_pb], + "read_options": read_options, + } + _set_database_id_to_request(expected_request_1, database_id) ds_api.lookup.assert_any_call( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key2_pb], - "read_options": read_options, - }, + request=expected_request_1, ) + expected_request_2 = { + "project_id": PROJECT, + "keys": [key1_pb, key2_pb], + "read_options": read_options, + } + _set_database_id_to_request(expected_request_2, database_id) ds_api.lookup.assert_any_call( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key1_pb, key2_pb], - "read_options": read_options, - }, + request=expected_request_2, ) -def test_client_get_multi_hit_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_hit_w_retry_w_timeout(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore.key import Key @@ -605,7 +629,7 @@ def test_client_get_multi_hit_w_retry_w_timeout(): timeout = 100000 # Make a found entity pb to be returned from mock backend. - entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo") + entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo", database=database_id) # Make a connection to return the entity pb. creds = _make_credentials() @@ -639,7 +663,8 @@ def test_client_get_multi_hit_w_retry_w_timeout(): ) -def test_client_get_multi_hit_w_transaction(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_hit_w_transaction(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore.key import Key @@ -649,16 +674,16 @@ def test_client_get_multi_hit_w_transaction(): path = [{"kind": kind, "id": id_}] # Make a found entity pb to be returned from mock backend. - entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo") + entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo", database=database_id) # Make a connection to return the entity pb. creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) lookup_response = _make_lookup_response(results=[entity_pb]) ds_api = _make_datastore_api(lookup_response=lookup_response) client._datastore_api_internal = ds_api - key = Key(kind, id_, project=PROJECT) + key = Key(kind, id_, project=PROJECT, database=database_id) txn = client.transaction() txn._id = txn_id (result,) = client.get_multi([key], transaction=txn) @@ -672,17 +697,17 @@ def test_client_get_multi_hit_w_transaction(): assert result["foo"] == "Foo" read_options = datastore_pb2.ReadOptions(transaction=txn_id) - ds_api.lookup.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key.to_protobuf()], - "read_options": read_options, - } - ) + expected_request = { + "project_id": PROJECT, + "keys": [key.to_protobuf()], + "read_options": read_options, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.lookup.assert_called_once_with(request=expected_request) -def test_client_get_multi_hit_w_read_time(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_hit_w_read_time(database_id): from datetime import datetime from google.cloud.datastore.key import Key @@ -696,16 +721,16 @@ def test_client_get_multi_hit_w_read_time(): path = [{"kind": kind, "id": id_}] # Make a found entity pb to be returned from mock backend. - entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo") + entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo", database=database_id) # Make a connection to return the entity pb. creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) lookup_response = _make_lookup_response(results=[entity_pb]) ds_api = _make_datastore_api(lookup_response=lookup_response) client._datastore_api_internal = ds_api - key = Key(kind, id_, project=PROJECT) + key = Key(kind, id_, project=PROJECT, database=database_id) (result,) = client.get_multi([key], read_time=read_time) new_key = result.key @@ -717,17 +742,17 @@ def test_client_get_multi_hit_w_read_time(): assert result["foo"] == "Foo" read_options = datastore_pb2.ReadOptions(read_time=read_time_pb) - ds_api.lookup.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key.to_protobuf()], - "read_options": read_options, - } - ) + expected_request = { + "project_id": PROJECT, + "keys": [key.to_protobuf()], + "read_options": read_options, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.lookup.assert_called_once_with(request=expected_request) -def test_client_get_multi_hit_multiple_keys_same_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_hit_multiple_keys_same_project(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore.key import Key @@ -736,18 +761,18 @@ def test_client_get_multi_hit_multiple_keys_same_project(): id2 = 2345 # Make a found entity pb to be returned from mock backend. - entity_pb1 = _make_entity_pb(PROJECT, kind, id1) - entity_pb2 = _make_entity_pb(PROJECT, kind, id2) + entity_pb1 = _make_entity_pb(PROJECT, kind, id1, database=database_id) + entity_pb2 = _make_entity_pb(PROJECT, kind, id2, database=database_id) # Make a connection to return the entity pbs. creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) lookup_response = _make_lookup_response(results=[entity_pb1, entity_pb2]) ds_api = _make_datastore_api(lookup_response=lookup_response) client._datastore_api_internal = ds_api - key1 = Key(kind, id1, project=PROJECT) - key2 = Key(kind, id2, project=PROJECT) + key1 = Key(kind, id1, project=PROJECT, database=database_id) + key2 = Key(kind, id2, project=PROJECT, database=database_id) retrieved1, retrieved2 = client.get_multi([key1, key2]) # Check values match. @@ -757,49 +782,50 @@ def test_client_get_multi_hit_multiple_keys_same_project(): assert dict(retrieved2) == {} read_options = datastore_pb2.ReadOptions() - ds_api.lookup.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "keys": [key1.to_protobuf(), key2.to_protobuf()], - "read_options": read_options, - } - ) + expected_request = { + "project_id": PROJECT, + "keys": [key1.to_protobuf(), key2.to_protobuf()], + "read_options": read_options, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.lookup.assert_called_once_with(request=expected_request) -def test_client_get_multi_hit_multiple_keys_different_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_hit_multiple_keys_different_project(database_id): from google.cloud.datastore.key import Key PROJECT1 = "PROJECT" PROJECT2 = "PROJECT-ALT" - key1 = Key("KIND", 1234, project=PROJECT1) - key2 = Key("KIND", 1234, project=PROJECT2) + key1 = Key("KIND", 1234, project=PROJECT1, database=database_id) + key2 = Key("KIND", 1234, project=PROJECT2, database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): client.get_multi([key1, key2]) -def test_client_get_multi_max_loops(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_get_multi_max_loops(database_id): from google.cloud.datastore.key import Key kind = "Kind" id_ = 1234 # Make a found entity pb to be returned from mock backend. - entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo") + entity_pb = _make_entity_pb(PROJECT, kind, id_, "foo", "Foo", database=database_id) # Make a connection to return the entity pb. creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) lookup_response = _make_lookup_response(results=[entity_pb]) ds_api = _make_datastore_api(lookup_response=lookup_response) client._datastore_api_internal = ds_api - key = Key(kind, id_, project=PROJECT) + key = Key(kind, id_, project=PROJECT, database=database_id) deferred = [] missing = [] @@ -815,10 +841,11 @@ def test_client_get_multi_max_loops(): ds_api.lookup.assert_not_called() -def test_client_put(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_put(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) put_multi = client.put_multi = mock.Mock() entity = mock.Mock() @@ -827,10 +854,11 @@ def test_client_put(): put_multi.assert_called_once_with(entities=[entity], retry=None, timeout=None) -def test_client_put_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_put_w_retry_w_timeout(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) put_multi = client.put_multi = mock.Mock() entity = mock.Mock() retry = mock.Mock() @@ -841,32 +869,35 @@ def test_client_put_w_retry_w_timeout(): put_multi.assert_called_once_with(entities=[entity], retry=retry, timeout=timeout) -def test_client_put_multi_no_entities(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_put_multi_no_entities(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) assert client.put_multi([]) is None -def test_client_put_multi_w_single_empty_entity(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_put_multi_w_single_empty_entity(database_id): # https://github.com/GoogleCloudPlatform/google-cloud-python/issues/649 from google.cloud.datastore.entity import Entity creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): client.put_multi(Entity()) -def test_client_put_multi_no_batch_w_partial_key_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_put_multi_no_batch_w_partial_key_w_retry_w_timeout(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 entity = _Entity(foo="bar") - key = entity.key = _Key(_Key.kind, None) + key = entity.key = _Key(_Key.kind, None, database=database_id) retry = mock.Mock() timeout = 100000 creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) key_pb = _make_key(234) ds_api = _make_datastore_api(key_pb) client._datastore_api_internal = ds_api @@ -874,14 +905,15 @@ def test_client_put_multi_no_batch_w_partial_key_w_retry_w_timeout(): result = client.put_multi([entity], retry=retry, timeout=timeout) assert result is None + expected_request = { + "project_id": PROJECT, + "mode": datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL, + "mutations": mock.ANY, + "transaction": None, + } + _set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "mode": datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL, - "mutations": mock.ANY, - "transaction": None, - }, + request=expected_request, retry=retry, timeout=timeout, ) @@ -897,11 +929,12 @@ def test_client_put_multi_no_batch_w_partial_key_w_retry_w_timeout(): assert value_pb.string_value == "bar" -def test_client_put_multi_existing_batch_w_completed_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_put_multi_existing_batch_w_completed_key(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) entity = _Entity(foo="bar") - key = entity.key = _Key() + key = entity.key = _Key(database=database_id) with _NoCommitBatch(client) as CURR_BATCH: result = client.put_multi([entity]) @@ -941,9 +974,10 @@ def test_client_delete_w_retry_w_timeout(): delete_multi.assert_called_once_with(keys=[key], retry=retry, timeout=timeout) -def test_client_delete_multi_no_keys(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_delete_multi_no_keys(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) client._datastore_api_internal = _make_datastore_api() result = client.delete_multi([]) @@ -951,29 +985,31 @@ def test_client_delete_multi_no_keys(): client._datastore_api_internal.commit.assert_not_called() -def test_client_delete_multi_no_batch_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_delete_multi_no_batch_w_retry_w_timeout(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 - key = _Key() + key = _Key(database=database_id) retry = mock.Mock() timeout = 100000 creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) ds_api = _make_datastore_api() client._datastore_api_internal = ds_api result = client.delete_multi([key], retry=retry, timeout=timeout) assert result is None + expected_request = { + "project_id": PROJECT, + "mode": datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL, + "mutations": mock.ANY, + "transaction": None, + } + _set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with( - request={ - "project_id": PROJECT, - "database_id": "", - "mode": datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL, - "mutations": mock.ANY, - "transaction": None, - }, + request=expected_request, retry=retry, timeout=timeout, ) @@ -983,12 +1019,13 @@ def test_client_delete_multi_no_batch_w_retry_w_timeout(): assert mutated_key == key.to_protobuf() -def test_client_delete_multi_w_existing_batch(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_delete_multi_w_existing_batch(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) client._datastore_api_internal = _make_datastore_api() - key = _Key() + key = _Key(database=database_id) with _NoCommitBatch(client) as CURR_BATCH: result = client.delete_multi([key]) @@ -999,12 +1036,13 @@ def test_client_delete_multi_w_existing_batch(): client._datastore_api_internal.commit.assert_not_called() -def test_client_delete_multi_w_existing_transaction(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_delete_multi_w_existing_transaction(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) client._datastore_api_internal = _make_datastore_api() - key = _Key() + key = _Key(database=database_id) with _NoCommitTransaction(client) as CURR_XACT: result = client.delete_multi([key]) @@ -1015,14 +1053,15 @@ def test_client_delete_multi_w_existing_transaction(): client._datastore_api_internal.commit.assert_not_called() -def test_client_delete_multi_w_existing_transaction_entity(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_delete_multi_w_existing_transaction_entity(database_id): from google.cloud.datastore.entity import Entity creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) client._datastore_api_internal = _make_datastore_api() - key = _Key() + key = _Key(database=database_id) entity = Entity(key=key) with _NoCommitTransaction(client) as CURR_XACT: @@ -1034,22 +1073,24 @@ def test_client_delete_multi_w_existing_transaction_entity(): client._datastore_api_internal.commit.assert_not_called() -def test_client_allocate_ids_w_completed_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_allocate_ids_w_completed_key(database_id): creds = _make_credentials() client = _make_client(credentials=creds) - complete_key = _Key() + complete_key = _Key(database=database_id) with pytest.raises(ValueError): client.allocate_ids(complete_key, 2) -def test_client_allocate_ids_w_partial_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_allocate_ids_w_partial_key(database_id): num_ids = 2 - incomplete_key = _Key(_Key.kind, None) + incomplete_key = _Key(_Key.kind, None, database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) allocated = mock.Mock(keys=[_KeyPB(i) for i in range(num_ids)], spec=["keys"]) alloc_ids = mock.Mock(return_value=allocated, spec=[]) ds_api = mock.Mock(allocate_ids=alloc_ids, spec=["allocate_ids"]) @@ -1061,43 +1102,21 @@ def test_client_allocate_ids_w_partial_key(): assert [key.id for key in result] == list(range(num_ids)) expected_keys = [incomplete_key.to_protobuf()] * num_ids - alloc_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys} - ) + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) + alloc_ids.assert_called_once_with(request=expected_request) -def test_client_allocate_ids_w_partial_key_w_specified_database(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_allocate_ids_w_partial_key_w_retry_w_timeout(database_id): num_ids = 2 - incomplete_key = _Key(_Key.kind, None, database=DATABASE) - - creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) - allocated = mock.Mock(keys=[_KeyPB(i) for i in range(num_ids)], spec=["keys"]) - alloc_ids = mock.Mock(return_value=allocated, spec=[]) - ds_api = mock.Mock(allocate_ids=alloc_ids, spec=["allocate_ids"]) - client._datastore_api_internal = ds_api - - result = client.allocate_ids(incomplete_key, num_ids) - - # Check the IDs returned. - assert [key.id for key in result] == list(range(num_ids)) - - expected_keys = [incomplete_key.to_protobuf()] * num_ids - alloc_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": DATABASE, "keys": expected_keys} - ) - - -def test_client_allocate_ids_w_partial_key_w_retry_w_timeout(): - num_ids = 2 - - incomplete_key = _Key(_Key.kind, None) + incomplete_key = _Key(_Key.kind, None, database=database_id) retry = mock.Mock() timeout = 100000 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) allocated = mock.Mock(keys=[_KeyPB(i) for i in range(num_ids)], spec=["keys"]) alloc_ids = mock.Mock(return_value=allocated, spec=[]) ds_api = mock.Mock(allocate_ids=alloc_ids, spec=["allocate_ids"]) @@ -1109,17 +1128,20 @@ def test_client_allocate_ids_w_partial_key_w_retry_w_timeout(): assert [key.id for key in result] == list(range(num_ids)) expected_keys = [incomplete_key.to_protobuf()] * num_ids + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) alloc_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys}, + request=expected_request, retry=retry, timeout=timeout, ) -def test_client_reserve_ids_sequential_w_completed_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_sequential_w_completed_key(database_id): num_ids = 2 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) complete_key = _Key() reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) @@ -1132,19 +1154,20 @@ def test_client_reserve_ids_sequential_w_completed_key(): _Key(_Key.kind, id) for id in range(complete_key.id, complete_key.id + num_ids) ) expected_keys = [key.to_protobuf() for key in reserved_keys] - reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys} - ) + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) + reserve_ids.assert_called_once_with(request=expected_request) -def test_client_reserve_ids_sequential_w_completed_key_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_sequential_w_completed_key_w_retry_w_timeout(database_id): num_ids = 2 retry = mock.Mock() timeout = 100000 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) - complete_key = _Key() + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) + complete_key = _Key(database=database_id) assert not complete_key.is_partial reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) @@ -1156,17 +1179,20 @@ def test_client_reserve_ids_sequential_w_completed_key_w_retry_w_timeout(): _Key(_Key.kind, id) for id in range(complete_key.id, complete_key.id + num_ids) ) expected_keys = [key.to_protobuf() for key in reserved_keys] + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys}, + request=expected_request, retry=retry, timeout=timeout, ) -def test_client_reserve_ids_sequential_w_completed_key_w_ancestor(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_sequential_w_completed_key_w_ancestor(database_id): num_ids = 2 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) complete_key = _Key("PARENT", "SINGLETON", _Key.kind, 1234) reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) @@ -1176,38 +1202,41 @@ def test_client_reserve_ids_sequential_w_completed_key_w_ancestor(): client.reserve_ids_sequential(complete_key, num_ids) reserved_keys = ( - _Key("PARENT", "SINGLETON", _Key.kind, id) + _Key("PARENT", "SINGLETON", _Key.kind, id, database=database_id) for id in range(complete_key.id, complete_key.id + num_ids) ) expected_keys = [key.to_protobuf() for key in reserved_keys] - reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys} - ) + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) + reserve_ids.assert_called_once_with(request=expected_request) -def test_client_reserve_ids_sequential_w_partial_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_sequential_w_partial_key(database_id): num_ids = 2 - incomplete_key = _Key(_Key.kind, None) + incomplete_key = _Key(_Key.kind, None, database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): client.reserve_ids_sequential(incomplete_key, num_ids) -def test_client_reserve_ids_sequential_w_wrong_num_ids(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_sequential_w_wrong_num_ids(database_id): num_ids = "2" - complete_key = _Key() + complete_key = _Key(database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): client.reserve_ids_sequential(complete_key, num_ids) -def test_client_reserve_ids_sequential_w_non_numeric_key_name(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_sequential_w_non_numeric_key_name(database_id): num_ids = 2 - complete_key = _Key(_Key.kind, "batman") + complete_key = _Key(_Key.kind, "batman", database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): client.reserve_ids_sequential(complete_key, num_ids) @@ -1217,13 +1246,14 @@ def _assert_reserve_ids_warning(warned): assert "Client.reserve_ids is deprecated." in str(warned[0].message) -def test_client_reserve_ids_w_partial_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_w_partial_key(database_id): import warnings num_ids = 2 incomplete_key = _Key(_Key.kind, None) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): with warnings.catch_warnings(record=True) as warned: client.reserve_ids(incomplete_key, num_ids) @@ -1231,13 +1261,14 @@ def test_client_reserve_ids_w_partial_key(): _assert_reserve_ids_warning(warned) -def test_client_reserve_ids_w_wrong_num_ids(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_w_wrong_num_ids(database_id): import warnings num_ids = "2" - complete_key = _Key() + complete_key = _Key(database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): with warnings.catch_warnings(record=True) as warned: client.reserve_ids(complete_key, num_ids) @@ -1245,13 +1276,14 @@ def test_client_reserve_ids_w_wrong_num_ids(): _assert_reserve_ids_warning(warned) -def test_client_reserve_ids_w_non_numeric_key_name(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_w_non_numeric_key_name(database_id): import warnings num_ids = 2 - complete_key = _Key(_Key.kind, "batman") + complete_key = _Key(_Key.kind, "batman", database=database_id) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): with warnings.catch_warnings(record=True) as warned: client.reserve_ids(complete_key, num_ids) @@ -1259,13 +1291,14 @@ def test_client_reserve_ids_w_non_numeric_key_name(): _assert_reserve_ids_warning(warned) -def test_client_reserve_ids_w_completed_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_w_completed_key(database_id): import warnings num_ids = 2 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) - complete_key = _Key() + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) + complete_key = _Key(database=database_id) reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) client._datastore_api_internal = ds_api @@ -1275,16 +1308,18 @@ def test_client_reserve_ids_w_completed_key(): client.reserve_ids(complete_key, num_ids) reserved_keys = ( - _Key(_Key.kind, id) for id in range(complete_key.id, complete_key.id + num_ids) + _Key(_Key.kind, id, database=database_id) + for id in range(complete_key.id, complete_key.id + num_ids) ) expected_keys = [key.to_protobuf() for key in reserved_keys] - reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys} - ) + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) + reserve_ids.assert_called_once_with(request=expected_request) _assert_reserve_ids_warning(warned) -def test_client_reserve_ids_w_completed_key_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_w_completed_key_w_retry_w_timeout(database_id): import warnings num_ids = 2 @@ -1292,8 +1327,8 @@ def test_client_reserve_ids_w_completed_key_w_retry_w_timeout(): timeout = 100000 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) - complete_key = _Key() + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) + complete_key = _Key(database=database_id) assert not complete_key.is_partial reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) @@ -1303,24 +1338,28 @@ def test_client_reserve_ids_w_completed_key_w_retry_w_timeout(): client.reserve_ids(complete_key, num_ids, retry=retry, timeout=timeout) reserved_keys = ( - _Key(_Key.kind, id) for id in range(complete_key.id, complete_key.id + num_ids) + _Key(_Key.kind, id, database=database_id) + for id in range(complete_key.id, complete_key.id + num_ids) ) expected_keys = [key.to_protobuf() for key in reserved_keys] + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys}, + request=expected_request, retry=retry, timeout=timeout, ) _assert_reserve_ids_warning(warned) -def test_client_reserve_ids_w_completed_key_w_ancestor(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_w_completed_key_w_ancestor(database_id): import warnings num_ids = 2 creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) - complete_key = _Key("PARENT", "SINGLETON", _Key.kind, 1234) + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) + complete_key = _Key("PARENT", "SINGLETON", _Key.kind, 1234, database=database_id) reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) client._datastore_api_internal = ds_api @@ -1330,41 +1369,44 @@ def test_client_reserve_ids_w_completed_key_w_ancestor(): client.reserve_ids(complete_key, num_ids) reserved_keys = ( - _Key("PARENT", "SINGLETON", _Key.kind, id) + _Key("PARENT", "SINGLETON", _Key.kind, id, database=database_id) for id in range(complete_key.id, complete_key.id + num_ids) ) expected_keys = [key.to_protobuf() for key in reserved_keys] - reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys} - ) + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) + + reserve_ids.assert_called_once_with(request=expected_request) _assert_reserve_ids_warning(warned) -def test_client_key_w_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_key_w_project(database_id): KIND = "KIND" ID = 1234 creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(TypeError): - client.key(KIND, ID, project=PROJECT) + client.key(KIND, ID, project=PROJECT, database=database_id) -def test_client_key_wo_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_key_wo_project(database_id): kind = "KIND" id_ = 1234 creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Key", spec=["__call__"]) with patch as mock_klass: key = client.key(kind, id_) assert key is mock_klass.return_value mock_klass.assert_called_once_with( - kind, id_, project=PROJECT, namespace=None, database="" + kind, id_, project=PROJECT, namespace=None, database=database_id ) @@ -1376,7 +1418,7 @@ def test_client_key_w_database(): client = _make_client(credentials=creds) with pytest.raises(TypeError): - client.key(KIND, ID, database=DATABASE) + client.key(KIND, ID, database="somedb") def test_client_key_wo_database(): @@ -1396,44 +1438,47 @@ def test_client_key_wo_database(): ) -def test_client_key_w_namespace(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_key_w_namespace(database_id): kind = "KIND" id_ = 1234 namespace = object() creds = _make_credentials() - client = _make_client(namespace=namespace, credentials=creds) + client = _make_client(namespace=namespace, credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Key", spec=["__call__"]) with patch as mock_klass: key = client.key(kind, id_) assert key is mock_klass.return_value mock_klass.assert_called_once_with( - kind, id_, project=PROJECT, namespace=namespace, database="" + kind, id_, project=PROJECT, namespace=namespace, database=database_id ) -def test_client_key_w_namespace_collision(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_key_w_namespace_collision(database_id): kind = "KIND" id_ = 1234 namespace1 = object() namespace2 = object() creds = _make_credentials() - client = _make_client(namespace=namespace1, credentials=creds) + client = _make_client(namespace=namespace1, credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Key", spec=["__call__"]) with patch as mock_klass: key = client.key(kind, id_, namespace=namespace2) assert key is mock_klass.return_value mock_klass.assert_called_once_with( - kind, id_, project=PROJECT, namespace=namespace2, database="" + kind, id_, project=PROJECT, namespace=namespace2, database=database_id ) -def test_client_entity_w_defaults(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_entity_w_defaults(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Entity", spec=["__call__"]) with patch as mock_klass: @@ -1503,40 +1548,31 @@ def test_client_query_w_other_client(): client.query(kind=KIND, client=other) -def test_client_query_w_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_query_w_project(database_id): KIND = "KIND" creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(TypeError): client.query(kind=KIND, project=PROJECT) -def test_client_query_w_database(): - KIND = "KIND" - +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_query_w_defaults(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) - - with pytest.raises(TypeError): - client.query(kind=KIND, database=DATABASE) - - -def test_client_query_w_defaults(): - creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Query", spec=["__call__"]) with patch as mock_klass: query = client.query() assert query is mock_klass.return_value - mock_klass.assert_called_once_with( - client, project=PROJECT, namespace=None, database="" - ) + mock_klass.assert_called_once_with(client, project=PROJECT, namespace=None) -def test_client_query_w_explicit(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_query_w_explicit(database_id): kind = "KIND" namespace = "NAMESPACE" ancestor = object() @@ -1546,7 +1582,7 @@ def test_client_query_w_explicit(): distinct_on = ["DISTINCT_ON"] creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Query", spec=["__call__"]) with patch as mock_klass: @@ -1565,7 +1601,6 @@ def test_client_query_w_explicit(): project=PROJECT, kind=kind, namespace=namespace, - database="", ancestor=ancestor, filters=filters, projection=projection, @@ -1574,42 +1609,45 @@ def test_client_query_w_explicit(): ) -def test_client_query_w_namespace(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_query_w_namespace(database_id): kind = "KIND" namespace = object() creds = _make_credentials() - client = _make_client(namespace=namespace, credentials=creds) + client = _make_client(namespace=namespace, credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Query", spec=["__call__"]) with patch as mock_klass: query = client.query(kind=kind) assert query is mock_klass.return_value mock_klass.assert_called_once_with( - client, project=PROJECT, namespace=namespace, kind=kind, database="" + client, project=PROJECT, namespace=namespace, kind=kind ) -def test_client_query_w_namespace_collision(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_query_w_namespace_collision(database_id): kind = "KIND" namespace1 = object() namespace2 = object() creds = _make_credentials() - client = _make_client(namespace=namespace1, credentials=creds) + client = _make_client(namespace=namespace1, credentials=creds, database=database_id) patch = mock.patch("google.cloud.datastore.client.Query", spec=["__call__"]) with patch as mock_klass: query = client.query(kind=kind, namespace=namespace2) assert query is mock_klass.return_value mock_klass.assert_called_once_with( - client, project=PROJECT, namespace=namespace2, kind=kind, database="" + client, project=PROJECT, namespace=namespace2, kind=kind ) -def test_client_aggregation_query_w_defaults(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_aggregation_query_w_defaults(database_id): creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) query = client.query() patch = mock.patch( "google.cloud.datastore.client.AggregationQuery", spec=["__call__"] @@ -1620,42 +1658,46 @@ def test_client_aggregation_query_w_defaults(): mock_klass.assert_called_once_with(client, query) -def test_client_aggregation_query_w_namespace(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_aggregation_query_w_namespace(database_id): namespace = object() creds = _make_credentials() - client = _make_client(namespace=namespace, credentials=creds) + client = _make_client(namespace=namespace, credentials=creds, database=database_id) query = client.query() aggregation_query = client.aggregation_query(query=query) assert aggregation_query.namespace == namespace -def test_client_aggregation_query_w_namespace_collision(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_aggregation_query_w_namespace_collision(database_id): namespace1 = object() namespace2 = object() creds = _make_credentials() - client = _make_client(namespace=namespace1, credentials=creds) + client = _make_client(namespace=namespace1, credentials=creds, database=database_id) query = client.query(namespace=namespace2) aggregation_query = client.aggregation_query(query=query) assert aggregation_query.namespace == namespace2 -def test_client_reserve_ids_multi_w_partial_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_multi_w_partial_key(database_id): incomplete_key = _Key(_Key.kind, None) creds = _make_credentials() - client = _make_client(credentials=creds) + client = _make_client(credentials=creds, database=database_id) with pytest.raises(ValueError): client.reserve_ids_multi([incomplete_key]) -def test_client_reserve_ids_multi(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_client_reserve_ids_multi(database_id): creds = _make_credentials() - client = _make_client(credentials=creds, _use_grpc=False) - key1 = _Key(_Key.kind, "one") - key2 = _Key(_Key.kind, "two") + client = _make_client(credentials=creds, _use_grpc=False, database=database_id) + key1 = _Key(_Key.kind, "one", database=database_id) + key2 = _Key(_Key.kind, "two", database=database_id) reserve_ids = mock.Mock() ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"]) client._datastore_api_internal = ds_api @@ -1663,9 +1705,9 @@ def test_client_reserve_ids_multi(): client.reserve_ids_multi([key1, key2]) expected_keys = [key1.to_protobuf(), key2.to_protobuf()] - reserve_ids.assert_called_once_with( - request={"project_id": PROJECT, "database_id": "", "keys": expected_keys} - ) + expected_request = {"project_id": PROJECT, "keys": expected_keys} + _set_database_id_to_request(expected_request, database_id) + reserve_ids.assert_called_once_with(request=expected_request) class _NoCommitBatch(object): @@ -1713,7 +1755,7 @@ class _Key(object): id = 1234 name = None _project = project = PROJECT - _database = database = "" + _database = database = None _namespace = None _key = "KEY" @@ -1838,12 +1880,13 @@ def _make_credentials(): return mock.Mock(spec=google.auth.credentials.Credentials) -def _make_entity_pb(project, kind, integer_id, name=None, str_val=None): +def _make_entity_pb(project, kind, integer_id, name=None, str_val=None, database=None): from google.cloud.datastore_v1.types import entity as entity_pb2 from google.cloud.datastore.helpers import _new_value_pb entity_pb = entity_pb2.Entity() entity_pb.key.partition_id.project_id = project + entity_pb.key.partition_id.database_id = database path_element = entity_pb._pb.key.path.add() path_element.kind = kind path_element.id = integer_id diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py index b2465719..467a2df1 100644 --- a/tests/unit/test_helpers.py +++ b/tests/unit/test_helpers.py @@ -455,42 +455,38 @@ def _make_key_pb(project=None, namespace=None, path=(), database=None): return pb -def test_key_from_protobuf_wo_database_or_namespace_in_pb(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_from_protobuf_wo_database_or_namespace_in_pb(database_id): from google.cloud.datastore.helpers import key_from_protobuf _PROJECT = "PROJECT" - pb = _make_key_pb(path=[{"kind": "KIND"}], project=_PROJECT) + pb = _make_key_pb(path=[{"kind": "KIND"}], project=_PROJECT, database=database_id) key = key_from_protobuf(pb) assert key.project == _PROJECT - assert key.database == "" + assert key.database == database_id assert key.namespace is None -def test_key_from_protobuf_w_database_in_pb(): - from google.cloud.datastore.helpers import key_from_protobuf - - _PROJECT = "PROJECT" - _DATABASE = "DATABASE" - pb = _make_key_pb(path=[{"kind": "KIND"}], project=_PROJECT, database=_DATABASE) - key = key_from_protobuf(pb) - assert key.project == _PROJECT - assert key.database == _DATABASE - assert key.namespace is None - - -def test_key_from_protobuf_w_namespace_in_pb(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_from_protobuf_w_namespace_in_pb(database_id): from google.cloud.datastore.helpers import key_from_protobuf _PROJECT = "PROJECT" _NAMESPACE = "NAMESPACE" - pb = _make_key_pb(path=[{"kind": "KIND"}], namespace=_NAMESPACE, project=_PROJECT) + pb = _make_key_pb( + path=[{"kind": "KIND"}], + namespace=_NAMESPACE, + project=_PROJECT, + database=database_id, + ) key = key_from_protobuf(pb) assert key.project == _PROJECT - assert key.database == "" + assert key.database == database_id assert key.namespace == _NAMESPACE -def test_key_from_protobuf_w_nested_path_in_pb(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_from_protobuf_w_nested_path_in_pb(database_id): from google.cloud.datastore.helpers import key_from_protobuf _PATH = [ @@ -498,9 +494,10 @@ def test_key_from_protobuf_w_nested_path_in_pb(): {"kind": "CHILD", "id": 1234}, {"kind": "GRANDCHILD", "id": 5678}, ] - pb = _make_key_pb(path=_PATH, project="PROJECT") + pb = _make_key_pb(path=_PATH, project="PROJECT", database=database_id) key = key_from_protobuf(pb) assert key.path == _PATH + assert key.database == database_id def test_w_nothing_in_pb(): diff --git a/tests/unit/test_key.py b/tests/unit/test_key.py index f0ac6332..e78981f2 100644 --- a/tests/unit/test_key.py +++ b/tests/unit/test_key.py @@ -90,19 +90,30 @@ def test_key_ctor_parent(): assert key.parent is parent_key -def test_key_ctor_partial_parent(): - parent_key = _make_key("KIND", project=_DEFAULT_PROJECT) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_ctor_partial_parent(database_id): + parent_key = _make_key("KIND", project=_DEFAULT_PROJECT, database=database_id) with pytest.raises(ValueError): - _make_key("KIND2", 1234, parent=parent_key) + _make_key("KIND2", 1234, parent=parent_key, database=database_id) -def test_key_ctor_parent_bad_type(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_ctor_parent_bad_type(database_id): with pytest.raises(AttributeError): - _make_key("KIND2", 1234, parent=("KIND1", 1234), project=_DEFAULT_PROJECT) + _make_key( + "KIND2", + 1234, + parent=("KIND1", 1234), + project=_DEFAULT_PROJECT, + database=database_id, + ) -def test_key_ctor_parent_bad_namespace(): - parent_key = _make_key("KIND", 1234, namespace="FOO", project=_DEFAULT_PROJECT) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_ctor_parent_bad_namespace(database_id): + parent_key = _make_key( + "KIND", 1234, namespace="FOO", project=_DEFAULT_PROJECT, database=database_id + ) with pytest.raises(ValueError): _make_key( "KIND2", @@ -110,13 +121,15 @@ def test_key_ctor_parent_bad_namespace(): namespace="BAR", parent=parent_key, PROJECT=_DEFAULT_PROJECT, + database=database_id, ) -def test_key_ctor_parent_bad_project(): - parent_key = _make_key("KIND", 1234, project="FOO") +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_ctor_parent_bad_project(database_id): + parent_key = _make_key("KIND", 1234, project="FOO", database=database_id) with pytest.raises(ValueError): - _make_key("KIND2", 1234, parent=parent_key, project="BAR") + _make_key("KIND2", 1234, parent=parent_key, project="BAR", database=database_id) def test_key_ctor_parent_empty_path(): @@ -142,15 +155,18 @@ def test_key_ctor_explicit(): assert key.path == _PATH -def test_key_ctor_explicit_w_unspecified_database(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_ctor_explicit_w_unspecified_database(database_id): _PROJECT = "PROJECT-ALT" _NAMESPACE = "NAMESPACE" _KIND = "KIND" _ID = 1234 _PATH = [{"kind": _KIND, "id": _ID}] - key = _make_key(_KIND, _ID, namespace=_NAMESPACE, project=_PROJECT) + key = _make_key( + _KIND, _ID, namespace=_NAMESPACE, project=_PROJECT, database=database_id + ) assert key.project == _PROJECT - assert key.database == _DEFAULT_DATABASE + assert key.database == database_id assert key.namespace == _NAMESPACE assert key.kind == _KIND assert key.path == _PATH @@ -192,9 +208,9 @@ def test_key__clone(): assert clone.path == _PATH -def test_key__clone_with_parent(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key__clone_with_parent(database_id): _PROJECT = "PROJECT-ALT" - _DATABASE = "DATABASE" _NAMESPACE = "NAMESPACE" _KIND1 = "PARENT" _KIND2 = "KIND" @@ -203,91 +219,99 @@ def test_key__clone_with_parent(): _PATH = [{"kind": _KIND1, "id": _ID1}, {"kind": _KIND2, "id": _ID2}] parent = _make_key( - _KIND1, _ID1, namespace=_NAMESPACE, database=_DATABASE, project=_PROJECT + _KIND1, _ID1, namespace=_NAMESPACE, database=database_id, project=_PROJECT ) - key = _make_key(_KIND2, _ID2, parent=parent) + key = _make_key(_KIND2, _ID2, parent=parent, database=database_id) assert key.parent is parent clone = key._clone() assert clone.parent is key.parent assert clone.project == _PROJECT - assert clone.database == _DATABASE + assert clone.database == database_id assert clone.namespace == _NAMESPACE assert clone.path == _PATH -def test_key___eq_____ne___w_non_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___w_non_key(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _NAME = "one" - key = _make_key(_KIND, _NAME, project=_PROJECT) + key = _make_key(_KIND, _NAME, project=_PROJECT, database=database_id) assert not key == object() assert key != object() -def test_key___eq_____ne___two_incomplete_keys_same_kind(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___two_incomplete_keys_same_kind(database_id): _PROJECT = "PROJECT" _KIND = "KIND" - key1 = _make_key(_KIND, project=_PROJECT) - key2 = _make_key(_KIND, project=_PROJECT) + key1 = _make_key(_KIND, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, project=_PROJECT, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___incomplete_key_w_complete_key_same_kind(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___incomplete_key_w_complete_key_same_kind(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _ID = 1234 - key1 = _make_key(_KIND, project=_PROJECT) - key2 = _make_key(_KIND, _ID, project=_PROJECT) + key1 = _make_key(_KIND, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, _ID, project=_PROJECT, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___complete_key_w_incomplete_key_same_kind(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___complete_key_w_incomplete_key_same_kind(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _ID = 1234 - key1 = _make_key(_KIND, _ID, project=_PROJECT) - key2 = _make_key(_KIND, project=_PROJECT) + key1 = _make_key(_KIND, _ID, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, project=_PROJECT, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___same_kind_different_ids(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_different_ids(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _ID1 = 1234 _ID2 = 2345 - key1 = _make_key(_KIND, _ID1, project=_PROJECT) - key2 = _make_key(_KIND, _ID2, project=_PROJECT) + key1 = _make_key(_KIND, _ID1, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, _ID2, project=_PROJECT, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___same_kind_and_id(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_id(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _ID = 1234 - key1 = _make_key(_KIND, _ID, project=_PROJECT) - key2 = _make_key(_KIND, _ID, project=_PROJECT) + key1 = _make_key(_KIND, _ID, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, _ID, project=_PROJECT, database=database_id) assert key1 == key2 assert not key1 != key2 -def test_key___eq_____ne___same_kind_and_id_different_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_id_different_project(database_id): _PROJECT1 = "PROJECT1" _PROJECT2 = "PROJECT2" _KIND = "KIND" _ID = 1234 - key1 = _make_key(_KIND, _ID, project=_PROJECT1) - key2 = _make_key(_KIND, _ID, project=_PROJECT2) + key1 = _make_key(_KIND, _ID, project=_PROJECT1, database=database_id) + key2 = _make_key(_KIND, _ID, project=_PROJECT2, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___same_kind_and_id_different_database(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_id_different_database(database_id): _PROJECT = "PROJECT" _DATABASE1 = "DATABASE1" _DATABASE2 = "DATABASE2" @@ -296,9 +320,11 @@ def test_key___eq_____ne___same_kind_and_id_different_database(): key1 = _make_key(_KIND, _ID, project=_PROJECT, database=_DATABASE1) key2 = _make_key(_KIND, _ID, project=_PROJECT, database=_DATABASE2) key_with_explicit_default = _make_key( - _KIND, _ID, project=_PROJECT, database=_DEFAULT_DATABASE + _KIND, _ID, project=_PROJECT, database=database_id + ) + key_with_implicit_default = _make_key( + _KIND, _ID, project=_PROJECT, database=database_id ) - key_with_implicit_default = _make_key(_KIND, _ID, project=_PROJECT) assert not key1 == key2 assert key1 != key2 assert not key1 == key_with_explicit_default @@ -308,87 +334,105 @@ def test_key___eq_____ne___same_kind_and_id_different_database(): assert key_with_explicit_default == key_with_implicit_default -def test_key___eq_____ne___same_kind_and_id_different_namespace(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_id_different_namespace(database_id): _PROJECT = "PROJECT" _NAMESPACE1 = "NAMESPACE1" _NAMESPACE2 = "NAMESPACE2" _KIND = "KIND" _ID = 1234 - key1 = _make_key(_KIND, _ID, project=_PROJECT, namespace=_NAMESPACE1) - key2 = _make_key(_KIND, _ID, project=_PROJECT, namespace=_NAMESPACE2) + key1 = _make_key( + _KIND, _ID, project=_PROJECT, namespace=_NAMESPACE1, database=database_id + ) + key2 = _make_key( + _KIND, _ID, project=_PROJECT, namespace=_NAMESPACE2, database=database_id + ) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___same_kind_different_names(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_different_names(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _NAME1 = "one" _NAME2 = "two" - key1 = _make_key(_KIND, _NAME1, project=_PROJECT) - key2 = _make_key(_KIND, _NAME2, project=_PROJECT) + key1 = _make_key(_KIND, _NAME1, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, _NAME2, project=_PROJECT, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___same_kind_and_name(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_name(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _NAME = "one" - key1 = _make_key(_KIND, _NAME, project=_PROJECT) - key2 = _make_key(_KIND, _NAME, project=_PROJECT) + key1 = _make_key(_KIND, _NAME, project=_PROJECT, database=database_id) + key2 = _make_key(_KIND, _NAME, project=_PROJECT, database=database_id) assert key1 == key2 assert not key1 != key2 -def test_key___eq_____ne___same_kind_and_name_different_project(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_name_different_project(database_id): _PROJECT1 = "PROJECT1" _PROJECT2 = "PROJECT2" _KIND = "KIND" _NAME = "one" - key1 = _make_key(_KIND, _NAME, project=_PROJECT1) - key2 = _make_key(_KIND, _NAME, project=_PROJECT2) + key1 = _make_key(_KIND, _NAME, project=_PROJECT1, database=database_id) + key2 = _make_key(_KIND, _NAME, project=_PROJECT2, database=database_id) assert not key1 == key2 assert key1 != key2 -def test_key___eq_____ne___same_kind_and_name_different_namespace(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___eq_____ne___same_kind_and_name_different_namespace(database_id): _PROJECT = "PROJECT" _NAMESPACE1 = "NAMESPACE1" _NAMESPACE2 = "NAMESPACE2" _KIND = "KIND" _NAME = "one" - key1 = _make_key(_KIND, _NAME, project=_PROJECT, namespace=_NAMESPACE1) - key2 = _make_key(_KIND, _NAME, project=_PROJECT, namespace=_NAMESPACE2) + key1 = _make_key( + _KIND, _NAME, project=_PROJECT, namespace=_NAMESPACE1, database=database_id + ) + key2 = _make_key( + _KIND, _NAME, project=_PROJECT, namespace=_NAMESPACE2, database=database_id + ) assert not key1 == key2 assert key1 != key2 -def test_key___hash___incomplete(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___hash___incomplete(database_id): _PROJECT = "PROJECT" _KIND = "KIND" - key = _make_key(_KIND, project=_PROJECT) - assert hash(key) != hash(_KIND) + hash(_PROJECT) + hash(None) + hash(None) + key = _make_key(_KIND, project=_PROJECT, database_id=database_id) + assert hash(key) != hash(_KIND) + hash(_PROJECT) + hash(None) + hash(None) + hash( + database_id + ) -def test_key___hash___completed_w_id(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___hash___completed_w_id(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _ID = 1234 - key = _make_key(_KIND, _ID, project=_PROJECT) + key = _make_key(_KIND, _ID, project=_PROJECT, database=database_id) assert hash(key) != hash(_KIND) + hash(_ID) + hash(_PROJECT) + hash(None) + hash( None - ) + ) + hash(database_id) -def test_key___hash___completed_w_name(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key___hash___completed_w_name(database_id): _PROJECT = "PROJECT" _KIND = "KIND" _NAME = "NAME" - key = _make_key(_KIND, _NAME, project=_PROJECT) + key = _make_key(_KIND, _NAME, project=_PROJECT, database=database_id) assert hash(key) != hash(_KIND) + hash(_NAME) + hash(_PROJECT) + hash(None) + hash( None - ) + ) + hash(database_id) def test_key___hash___completed_w_database_and_namespace(): @@ -405,13 +449,15 @@ def test_key___hash___completed_w_database_and_namespace(): ) + hash(None) -def test_key_completed_key_on_partial_w_id(): - key = _make_key("KIND", project=_DEFAULT_PROJECT) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_key_completed_key_on_partial_w_id(database_id): + key = _make_key("KIND", project=_DEFAULT_PROJECT, database=database_id) _ID = 1234 new_key = key.completed_key(_ID) assert key is not new_key assert new_key.id == _ID assert new_key.name is None + assert new_key.database == database_id def test_key_completed_key_on_partial_w_name(): @@ -547,7 +593,7 @@ def test_key_from_legacy_urlsafe(): key = Key.from_legacy_urlsafe(_URLSAFE_EXAMPLE1) assert "s~" + key.project == _URLSAFE_APP1 - assert key.database == _DEFAULT_DATABASE + assert key.database is None assert key.namespace == _URLSAFE_NAMESPACE1 assert key.flat_path == _URLSAFE_FLAT_PATH1 # Also make sure we didn't accidentally set the parent. diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py index ccf9f691..2b84c731 100644 --- a/tests/unit/test_query.py +++ b/tests/unit/test_query.py @@ -25,16 +25,19 @@ BaseCompositeFilter, ) +from google.cloud.datastore.helpers import _set_database_id_to_request + _PROJECT = "PROJECT" -def test_query_ctor_defaults(): - client = _make_client() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ctor_defaults(database_id): + client = _make_client(database=database_id) query = _make_query(client) assert query._client is client + assert query._client.database == client.database assert query.project == client.project assert query.kind is None - assert query.database == client.database assert query.namespace == client.namespace assert query.ancestor is None assert query.filters == [] @@ -52,15 +55,15 @@ def test_query_ctor_defaults(): [Or([PropertyFilter("foo", "=", "Qux"), PropertyFilter("bar", "<", 17)])], ], ) -def test_query_ctor_explicit(filters): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ctor_explicit(filters, database_id): from google.cloud.datastore.key import Key _PROJECT = "OTHER_PROJECT" - _DATABASE = "OTHER_DATABASE" _KIND = "KIND" _NAMESPACE = "OTHER_NAMESPACE" - client = _make_client() - ancestor = Key("ANCESTOR", 123, project=_PROJECT) + client = _make_client(database=database_id) + ancestor = Key("ANCESTOR", 123, project=_PROJECT, database=database_id) FILTERS = filters PROJECTION = ["foo", "bar", "baz"] ORDER = ["foo", "bar"] @@ -70,7 +73,6 @@ def test_query_ctor_explicit(filters): client, kind=_KIND, project=_PROJECT, - database=_DATABASE, namespace=_NAMESPACE, ancestor=ancestor, filters=FILTERS, @@ -79,8 +81,8 @@ def test_query_ctor_explicit(filters): distinct_on=DISTINCT_ON, ) assert query._client is client + assert query._client.database == database_id assert query.project == _PROJECT - assert query.database == _DATABASE assert query.kind == _KIND assert query.namespace == _NAMESPACE assert query.ancestor.path == ancestor.path @@ -90,88 +92,91 @@ def test_query_ctor_explicit(filters): assert query.distinct_on == DISTINCT_ON -def test_query_ctor_bad_projection(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ctor_bad_projection(database_id): BAD_PROJECTION = object() with pytest.raises(TypeError): - _make_query(_make_client(), projection=BAD_PROJECTION) + _make_query(_make_client(database=database_id), projection=BAD_PROJECTION) -def test_query_ctor_bad_order(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ctor_bad_order(database_id): BAD_ORDER = object() with pytest.raises(TypeError): - _make_query(_make_client(), order=BAD_ORDER) + _make_query(_make_client(database=database_id), order=BAD_ORDER) -def test_query_ctor_bad_distinct_on(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ctor_bad_distinct_on(database_id): BAD_DISTINCT_ON = object() with pytest.raises(TypeError): - _make_query(_make_client(), distinct_on=BAD_DISTINCT_ON) + _make_query(_make_client(database=database_id), distinct_on=BAD_DISTINCT_ON) -def test_query_ctor_bad_filters(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ctor_bad_filters(database_id): FILTERS_CANT_UNPACK = [("one", "two")] with pytest.raises(ValueError): - _make_query(_make_client(), filters=FILTERS_CANT_UNPACK) + _make_query(_make_client(database=database_id), filters=FILTERS_CANT_UNPACK) -def test_query_project_getter(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_project_getter(database_id): PROJECT = "PROJECT" - query = _make_query(_make_client(), project=PROJECT) + query = _make_query(_make_client(database=database_id), project=PROJECT) assert query.project == PROJECT -def test_query_database_getter(): - DATABASE = "DATABASE" - OTHER_DATABASE = "OTHER-DATABASE" - query = _make_query(_make_client(), database=DATABASE) - assert query.database == DATABASE - - # Fallback to client - client = _make_client() - client.database = None - query = _make_query(client) - client.database = OTHER_DATABASE - assert query.database == OTHER_DATABASE +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_database_getter(database_id): + query = _make_query(_make_client(database=database_id)) + assert query._client.database == database_id -def test_query_namespace_setter_w_non_string(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_namespace_setter_w_non_string(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(ValueError): query.namespace = object() -def test_query_namespace_setter(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_namespace_setter(database_id): _NAMESPACE = "OTHER_NAMESPACE" - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) query.namespace = _NAMESPACE assert query.namespace == _NAMESPACE -def test_query_kind_setter_w_non_string(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_kind_setter_w_non_string(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(TypeError): query.kind = object() -def test_query_kind_setter_wo_existing(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_kind_setter_wo_existing(database_id): _KIND = "KIND" - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) query.kind = _KIND assert query.kind == _KIND -def test_query_kind_setter_w_existing(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_kind_setter_w_existing(database_id): _KIND_BEFORE = "KIND_BEFORE" _KIND_AFTER = "KIND_AFTER" - query = _make_query(_make_client(), kind=_KIND_BEFORE) + query = _make_query(_make_client(database=database_id), kind=_KIND_BEFORE) assert query.kind == _KIND_BEFORE query.kind = _KIND_AFTER assert query.project == _PROJECT assert query.kind == _KIND_AFTER -def test_query_ancestor_setter_w_non_key(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ancestor_setter_w_non_key(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(TypeError): query.ancestor = object() @@ -180,68 +185,76 @@ def test_query_ancestor_setter_w_non_key(): query.ancestor = ["KIND", "NAME"] -def test_query_ancestor_setter_w_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ancestor_setter_w_key(database_id): from google.cloud.datastore.key import Key _NAME = "NAME" - key = Key("KIND", 123, project=_PROJECT) - query = _make_query(_make_client()) + key = Key("KIND", 123, project=_PROJECT, database=database_id) + query = _make_query(_make_client(database=database_id)) query.add_filter("name", "=", _NAME) query.ancestor = key assert query.ancestor.path == key.path -def test_query_ancestor_setter_w_key_property_filter(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ancestor_setter_w_key_property_filter(database_id): from google.cloud.datastore.key import Key _NAME = "NAME" - key = Key("KIND", 123, project=_PROJECT) - query = _make_query(_make_client()) + key = Key("KIND", 123, project=_PROJECT, database=database_id) + query = _make_query(_make_client(database=database_id)) query.add_filter(filter=PropertyFilter("name", "=", _NAME)) query.ancestor = key assert query.ancestor.path == key.path -def test_query_ancestor_deleter_w_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_ancestor_deleter_w_key(database_id): from google.cloud.datastore.key import Key - key = Key("KIND", 123, project=_PROJECT) - query = _make_query(client=_make_client(), ancestor=key) + key = Key("KIND", 123, project=_PROJECT, database=database_id) + query = _make_query(client=_make_client(database=database_id), ancestor=key) del query.ancestor assert query.ancestor is None -def test_query_add_filter_setter_w_unknown_operator(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_setter_w_unknown_operator(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(ValueError) as exc: query.add_filter("firstname", "~~", "John") assert "Invalid expression:" in str(exc.value) assert "Please use one of: =, <, <=, >, >=, !=, IN, NOT_IN." in str(exc.value) -def test_query_add_property_filter_setter_w_unknown_operator(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter_setter_w_unknown_operator(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(ValueError) as exc: query.add_filter(filter=PropertyFilter("firstname", "~~", "John")) assert "Invalid expression:" in str(exc.value) assert "Please use one of: =, <, <=, >, >=, !=, IN, NOT_IN." in str(exc.value) -def test_query_add_filter_w_known_operator(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_w_known_operator(database_id): + query = _make_query(_make_client(database=database_id)) query.add_filter("firstname", "=", "John") assert query.filters == [("firstname", "=", "John")] -def test_query_add_property_filter_w_known_operator(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter_w_known_operator(database_id): + query = _make_query(_make_client(database=database_id)) property_filter = PropertyFilter("firstname", "=", "John") query.add_filter(filter=property_filter) assert query.filters == [property_filter] -def test_query_add_filter_w_all_operators(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_w_all_operators(database_id): + query = _make_query(_make_client(database=database_id)) query.add_filter("leq_prop", "<=", "val1") query.add_filter("geq_prop", ">=", "val2") query.add_filter("lt_prop", "<", "val3") @@ -261,8 +274,9 @@ def test_query_add_filter_w_all_operators(): assert query.filters[7] == ("not_in_prop", "NOT_IN", ["val13"]) -def test_query_add_property_filter_w_all_operators(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter_w_all_operators(database_id): + query = _make_query(_make_client(database=database_id)) filters = [ ("leq_prop", "<=", "val1"), ("geq_prop", ">=", "val2"), @@ -284,10 +298,11 @@ def test_query_add_property_filter_w_all_operators(): assert query.filters[i] == property_filters[i] -def test_query_add_filter_w_known_operator_and_entity(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_w_known_operator_and_entity(database_id): from google.cloud.datastore.entity import Entity - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) other = Entity() other["firstname"] = "John" other["lastname"] = "Smith" @@ -295,10 +310,11 @@ def test_query_add_filter_w_known_operator_and_entity(): assert query.filters == [("other", "=", other)] -def test_query_add_property_filter_w_known_operator_and_entity(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter_w_known_operator_and_entity(database_id): from google.cloud.datastore.entity import Entity - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) other = Entity() other["firstname"] = "John" other["lastname"] = "Smith" @@ -307,52 +323,58 @@ def test_query_add_property_filter_w_known_operator_and_entity(): assert query.filters == [property_filter] -def test_query_add_filter_w_whitespace_property_name(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_w_whitespace_property_name(database_id): + query = _make_query(_make_client(database=database_id)) PROPERTY_NAME = " property with lots of space " query.add_filter(PROPERTY_NAME, "=", "John") assert query.filters == [(PROPERTY_NAME, "=", "John")] -def test_query_add_property_filter_w_whitespace_property_name(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter_w_whitespace_property_name(database_id): + query = _make_query(_make_client(database=database_id)) PROPERTY_NAME = " property with lots of space " property_filter = PropertyFilter(PROPERTY_NAME, "=", "John") query.add_filter(filter=property_filter) assert query.filters == [property_filter] -def test_query_add_filter___key__valid_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter___key__valid_key(database_id): from google.cloud.datastore.key import Key - query = _make_query(_make_client()) - key = Key("Foo", project=_PROJECT) + query = _make_query(_make_client(database=database_id)) + key = Key("Foo", project=_PROJECT, database=database_id) query.add_filter("__key__", "=", key) assert query.filters == [("__key__", "=", key)] -def test_query_add_property_filter___key__valid_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter___key__valid_key(database_id): from google.cloud.datastore.key import Key - query = _make_query(_make_client()) - key = Key("Foo", project=_PROJECT) + query = _make_query(_make_client(database=database_id)) + key = Key("Foo", project=_PROJECT, database=database_id) property_filter = PropertyFilter("__key__", "=", key) query.add_filter(filter=property_filter) assert query.filters == [property_filter] -def test_query_add_filter_return_query_obj(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_return_query_obj(database_id): from google.cloud.datastore.query import Query - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) query_obj = query.add_filter("firstname", "=", "John") assert isinstance(query_obj, Query) assert query_obj.filters == [("firstname", "=", "John")] -def test_query_add_property_filter_without_keyword_argument(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_property_filter_without_keyword_argument(database_id): - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) property_filter = PropertyFilter("firstname", "=", "John") with pytest.raises(ValueError) as exc: query.add_filter(property_filter) @@ -363,9 +385,10 @@ def test_query_add_property_filter_without_keyword_argument(): ) -def test_query_add_composite_filter_without_keyword_argument(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_composite_filter_without_keyword_argument(database_id): - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) and_filter = And(["firstname", "=", "John"]) with pytest.raises(ValueError) as exc: query.add_filter(and_filter) @@ -385,9 +408,10 @@ def test_query_add_composite_filter_without_keyword_argument(): ) -def test_query_positional_args_and_property_filter(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_positional_args_and_property_filter(database_id): - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) with pytest.raises(ValueError) as exc: query.add_filter("firstname", "=", "John", filter=("name", "=", "Blabla")) @@ -397,9 +421,10 @@ def test_query_positional_args_and_property_filter(): ) -def test_query_positional_args_and_composite_filter(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_positional_args_and_composite_filter(database_id): - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) and_filter = And(["firstname", "=", "John"]) with pytest.raises(ValueError) as exc: query.add_filter("firstname", "=", "John", filter=and_filter) @@ -410,8 +435,9 @@ def test_query_positional_args_and_composite_filter(): ) -def test_query_add_filter_with_positional_args_raises_user_warning(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_add_filter_with_positional_args_raises_user_warning(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.warns( UserWarning, match="Detected filter using positional arguments", @@ -425,151 +451,171 @@ def test_query_add_filter_with_positional_args_raises_user_warning(): _make_stub_query(filters=[("name", "=", "John")]) -def test_query_filter___key__not_equal_operator(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_filter___key__not_equal_operator(database_id): from google.cloud.datastore.key import Key - key = Key("Foo", project=_PROJECT) - query = _make_query(_make_client()) + key = Key("Foo", project=_PROJECT, database=database_id) + query = _make_query(_make_client(database=database_id)) query.add_filter("__key__", "<", key) assert query.filters == [("__key__", "<", key)] -def test_query_property_filter___key__not_equal_operator(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_property_filter___key__not_equal_operator(database_id): from google.cloud.datastore.key import Key - key = Key("Foo", project=_PROJECT) - query = _make_query(_make_client()) + key = Key("Foo", project=_PROJECT, database=database_id) + query = _make_query(_make_client(database=database_id)) property_filter = PropertyFilter("__key__", "<", key) query.add_filter(filter=property_filter) assert query.filters == [property_filter] -def test_query_filter___key__invalid_value(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_filter___key__invalid_value(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(ValueError) as exc: query.add_filter("__key__", "=", None) assert "Invalid key:" in str(exc.value) -def test_query_property_filter___key__invalid_value(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_property_filter___key__invalid_value(database_id): + query = _make_query(_make_client(database=database_id)) with pytest.raises(ValueError) as exc: query.add_filter(filter=PropertyFilter("__key__", "=", None)) assert "Invalid key:" in str(exc.value) -def test_query_projection_setter_empty(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_projection_setter_empty(database_id): + query = _make_query(_make_client(database=database_id)) query.projection = [] assert query.projection == [] -def test_query_projection_setter_string(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_projection_setter_string(database_id): + query = _make_query(_make_client(database=database_id)) query.projection = "field1" assert query.projection == ["field1"] -def test_query_projection_setter_non_empty(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_projection_setter_non_empty(database_id): + query = _make_query(_make_client(database=database_id)) query.projection = ["field1", "field2"] assert query.projection == ["field1", "field2"] -def test_query_projection_setter_multiple_calls(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_projection_setter_multiple_calls(database_id): _PROJECTION1 = ["field1", "field2"] _PROJECTION2 = ["field3"] - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) query.projection = _PROJECTION1 assert query.projection == _PROJECTION1 query.projection = _PROJECTION2 assert query.projection == _PROJECTION2 -def test_query_keys_only(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_keys_only(database_id): + query = _make_query(_make_client(database=database_id)) query.keys_only() assert query.projection == ["__key__"] -def test_query_key_filter_defaults(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_key_filter_defaults(database_id): from google.cloud.datastore.key import Key - client = _make_client() + client = _make_client(database=database_id) query = _make_query(client) assert query.filters == [] - key = Key("Kind", 1234, project="project") + key = Key("Kind", 1234, project="project", database=database_id) query.key_filter(key) assert query.filters == [("__key__", "=", key)] -def test_query_key_filter_explicit(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_key_filter_explicit(database_id): from google.cloud.datastore.key import Key - client = _make_client() + client = _make_client(database=database_id) query = _make_query(client) assert query.filters == [] - key = Key("Kind", 1234, project="project") + key = Key("Kind", 1234, project="project", database=database_id) query.key_filter(key, operator=">") assert query.filters == [("__key__", ">", key)] -def test_query_order_setter_empty(): - query = _make_query(_make_client(), order=["foo", "-bar"]) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_order_setter_empty(database_id): + query = _make_query(_make_client(database=database_id), order=["foo", "-bar"]) query.order = [] assert query.order == [] -def test_query_order_setter_string(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_order_setter_string(database_id): + query = _make_query(_make_client(database=database_id)) query.order = "field" assert query.order == ["field"] -def test_query_order_setter_single_item_list_desc(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_order_setter_single_item_list_desc(database_id): + query = _make_query(_make_client(database=database_id)) query.order = ["-field"] assert query.order == ["-field"] -def test_query_order_setter_multiple(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_order_setter_multiple(database_id): + query = _make_query(_make_client(database=database_id)) query.order = ["foo", "-bar"] assert query.order == ["foo", "-bar"] -def test_query_distinct_on_setter_empty(): - query = _make_query(_make_client(), distinct_on=["foo", "bar"]) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_distinct_on_setter_empty(database_id): + query = _make_query(_make_client(database=database_id), distinct_on=["foo", "bar"]) query.distinct_on = [] assert query.distinct_on == [] -def test_query_distinct_on_setter_string(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_distinct_on_setter_string(database_id): + query = _make_query(_make_client(database=database_id)) query.distinct_on = "field1" assert query.distinct_on == ["field1"] -def test_query_distinct_on_setter_non_empty(): - query = _make_query(_make_client()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_distinct_on_setter_non_empty(database_id): + query = _make_query(_make_client(database=database_id)) query.distinct_on = ["field1", "field2"] assert query.distinct_on == ["field1", "field2"] -def test_query_distinct_on_multiple_calls(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_distinct_on_multiple_calls(database_id): _DISTINCT_ON1 = ["field1", "field2"] _DISTINCT_ON2 = ["field3"] - query = _make_query(_make_client()) + query = _make_query(_make_client(database=database_id)) query.distinct_on = _DISTINCT_ON1 assert query.distinct_on == _DISTINCT_ON1 query.distinct_on = _DISTINCT_ON2 assert query.distinct_on == _DISTINCT_ON2 -def test_query_fetch_defaults_w_client_attr(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_fetch_defaults_w_client_attr(database_id): from google.cloud.datastore.query import Iterator - client = _make_client() + client = _make_client(database=database_id) query = _make_query(client) iterator = query.fetch() @@ -583,11 +629,12 @@ def test_query_fetch_defaults_w_client_attr(): assert iterator._timeout is None -def test_query_fetch_w_explicit_client_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_query_fetch_w_explicit_client_w_retry_w_timeout(database_id): from google.cloud.datastore.query import Iterator - client = _make_client() - other_client = _make_client() + client = _make_client(database=database_id) + other_client = _make_client(database=database_id) query = _make_query(client) retry = mock.Mock() timeout = 100000 @@ -721,13 +768,14 @@ def test_iterator__build_protobuf_all_values_except_start_and_end_cursor(): assert pb == expected_pb -def test_iterator__process_query_results(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__process_query_results(database_id): from google.cloud.datastore_v1.types import query as query_pb2 iterator = _make_iterator(None, None, end_cursor="abcd") assert iterator._end_cursor is not None - entity_pbs = [_make_entity("Hello", 9998, "PRAHJEKT")] + entity_pbs = [_make_entity("Hello", 9998, "PRAHJEKT", database=database_id)] cursor_as_bytes = b"\x9ai\xe7" cursor = b"mmnn" skipped_results = 4 @@ -743,13 +791,14 @@ def test_iterator__process_query_results(): assert iterator._more_results -def test_iterator__process_query_results_done(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__process_query_results_done(database_id): from google.cloud.datastore_v1.types import query as query_pb2 iterator = _make_iterator(None, None, end_cursor="abcd") assert iterator._end_cursor is not None - entity_pbs = [_make_entity("World", 1234, "PROJECT")] + entity_pbs = [_make_entity("World", 1234, "PROJECT", database=database_id)] cursor_as_bytes = b"\x9ai\xe7" skipped_results = 44 more_results_enum = query_pb2.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS @@ -774,7 +823,7 @@ def test_iterator__process_query_results_bad_enum(): def _next_page_helper( - txn_id=None, retry=None, timeout=None, read_time=None, database="" + txn_id=None, retry=None, timeout=None, read_time=None, database=None ): from google.api_core import page_iterator from google.cloud.datastore.query import Query @@ -825,45 +874,48 @@ def _next_page_helper( else: read_options = datastore_pb2.ReadOptions() empty_query = query_pb2.Query() + expected_request = { + "project_id": project, + "partition_id": partition_id, + "read_options": read_options, + "query": empty_query, + } + _set_database_id_to_request(expected_request, database) ds_api.run_query.assert_called_once_with( - request={ - "project_id": project, - "database_id": database, - "partition_id": partition_id, - "read_options": read_options, - "query": empty_query, - }, + request=expected_request, **kwargs, ) -def test_iterator__next_page(): - _next_page_helper() +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page(database_id): + _next_page_helper(database_id) -def test_iterator__next_page_w_retry(): - _next_page_helper(retry=mock.Mock()) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page_w_retry(database_id): + _next_page_helper(retry=mock.Mock(), database=database_id) -def test_iterator__next_page_w_timeout(): - _next_page_helper(timeout=100000) +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page_w_timeout(database_id): + _next_page_helper(timeout=100000, database=database_id) -def test_iterator__next_page_in_transaction(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page_in_transaction(database_id): txn_id = b"1xo1md\xe2\x98\x83" - _next_page_helper(txn_id) + _next_page_helper(txn_id, database=database_id) -def test_iterator__next_page_w_read_time(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page_w_read_time(database_id): read_time = datetime.datetime.utcfromtimestamp(1641058200.123456) - _next_page_helper(read_time=read_time) - + _next_page_helper(read_time=read_time, database=database_id) -def test_iterator__next_page_w_database(): - _next_page_helper(database="base-of-data") - -def test_iterator__next_page_no_more(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page_no_more(database_id): from google.cloud.datastore.query import Query ds_api = _make_datastore_api() @@ -877,7 +929,8 @@ def test_iterator__next_page_no_more(): ds_api.run_query.assert_not_called() -def test_iterator__next_page_w_skipped_lt_offset(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_iterator__next_page_w_skipped_lt_offset(database_id): from google.api_core import page_iterator from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 @@ -898,7 +951,7 @@ def test_iterator__next_page_w_skipped_lt_offset(): result_2.batch.skipped_cursor = skipped_cursor_2 ds_api = _make_datastore_api(result_1, result_2) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) query = Query(client) offset = 150 @@ -909,25 +962,24 @@ def test_iterator__next_page_w_skipped_lt_offset(): assert isinstance(page, page_iterator.Page) assert page._parent is iterator - partition_id = entity_pb2.PartitionId(project_id=project) + partition_id = entity_pb2.PartitionId(project_id=project, database_id=database_id) read_options = datastore_pb2.ReadOptions() query_1 = query_pb2.Query(offset=offset) query_2 = query_pb2.Query( start_cursor=skipped_cursor_1, offset=(offset - skipped_1) ) - expected_calls = [ - mock.call( - request={ - "project_id": project, - "database_id": "", - "partition_id": partition_id, - "read_options": read_options, - "query": query, - } - ) - for query in [query_1, query_2] - ] + expected_calls = [] + for query in [query_1, query_2]: + expected_request = { + "project_id": project, + "partition_id": partition_id, + "read_options": read_options, + "query": query, + } + _set_database_id_to_request(expected_request, database_id) + expected_calls.append(mock.call(request=expected_request)) + assert ds_api.run_query.call_args_list == expected_calls @@ -977,12 +1029,13 @@ def test_pb_from_query_kind(): assert [item.name for item in pb.kind] == ["KIND"] -def test_pb_from_query_ancestor(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_pb_from_query_ancestor(database_id): from google.cloud.datastore.key import Key from google.cloud.datastore_v1.types import query as query_pb2 from google.cloud.datastore.query import _pb_from_query - ancestor = Key("Ancestor", 123, project="PROJECT") + ancestor = Key("Ancestor", 123, project="PROJECT", database=database_id) pb = _pb_from_query(_make_stub_query(ancestor=ancestor)) cfilter = pb.filter.composite_filter assert cfilter.op == query_pb2.CompositeFilter.Operator.AND @@ -1008,12 +1061,13 @@ def test_pb_from_query_filter(): assert pfilter.value.string_value == "John" -def test_pb_from_query_filter_key(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_pb_from_query_filter_key(database_id): from google.cloud.datastore.key import Key from google.cloud.datastore_v1.types import query as query_pb2 from google.cloud.datastore.query import _pb_from_query - key = Key("Kind", 123, project="PROJECT") + key = Key("Kind", 123, project="PROJECT", database=database_id) query = _make_stub_query(filters=[("__key__", "=", key)]) query.OPERATORS = {"=": query_pb2.PropertyFilter.Operator.EQUAL} pb = _pb_from_query(query) @@ -1160,8 +1214,6 @@ def _make_stub_query( projection=(), order=(), distinct_on=(), - *, - database=None, ): query = Query( client, @@ -1173,7 +1225,6 @@ def _make_stub_query( projection=projection, order=order, distinct_on=distinct_on, - database=database ) return query @@ -1185,8 +1236,7 @@ def __init__( datastore_api=None, namespace=None, transaction=None, - *, - database="", + database=None, ): self.project = project self._datastore_api = datastore_api @@ -1211,15 +1261,16 @@ def _make_iterator(*args, **kw): return Iterator(*args, **kw) -def _make_client(): - return _Client(_PROJECT) +def _make_client(database=None): + return _Client(_PROJECT, database=database) -def _make_entity(kind, id_, project): +def _make_entity(kind, id_, project, database=None): from google.cloud.datastore_v1.types import entity as entity_pb2 key = entity_pb2.Key() key.partition_id.project_id = project + key.partition_id.database_id = database elem = key.path._pb.add() elem.kind = kind elem.id = id_ diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py index ce66eb4e..96854db6 100644 --- a/tests/unit/test_transaction.py +++ b/tests/unit/test_transaction.py @@ -15,18 +15,20 @@ import mock import pytest +from google.cloud.datastore.helpers import _set_database_id_to_request -def test_transaction_ctor_defaults(): + +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_ctor_defaults(database_id): from google.cloud.datastore.transaction import Transaction project = "PROJECT" - database = "DATABASE" - client = _Client(project, database=database) + client = _Client(project, database=database_id) xact = _make_transaction(client) assert xact.project == project - assert xact.database == database + assert xact.database == database_id assert xact._client is client assert xact.id is None assert xact._status == Transaction._INITIAL @@ -34,54 +36,59 @@ def test_transaction_ctor_defaults(): assert len(xact._partial_key_entities) == 0 -def test_transaction_constructor_read_only(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_constructor_read_only(database_id): project = "PROJECT" id_ = 850302 ds_api = _make_datastore_api(xact=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) options = _make_options(read_only=True) xact = _make_transaction(client, read_only=True) assert xact._options == options + assert xact.database == database_id -def test_transaction_constructor_w_read_time(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_constructor_w_read_time(database_id): from datetime import datetime project = "PROJECT" id_ = 850302 read_time = datetime.utcfromtimestamp(1641058200.123456) ds_api = _make_datastore_api(xact=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) options = _make_options(read_only=True, read_time=read_time) xact = _make_transaction(client, read_only=True, read_time=read_time) assert xact._options == options + assert xact.database == database_id -def test_transaction_constructor_read_write_w_read_time(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_constructor_read_write_w_read_time(database_id): from datetime import datetime project = "PROJECT" id_ = 850302 read_time = datetime.utcfromtimestamp(1641058200.123456) ds_api = _make_datastore_api(xact=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) with pytest.raises(ValueError): _make_transaction(client, read_only=False, read_time=read_time) -def test_transaction_current(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_current(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" - database = "DATABASE" id_ = 678 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, database=database, datastore_api=ds_api) + client = _Client(project, database=database_id, datastore_api=ds_api) xact1 = _make_transaction(client) xact2 = _make_transaction(client) assert xact1.current() is None @@ -111,89 +118,97 @@ def test_transaction_current(): begin_txn = ds_api.begin_transaction assert begin_txn.call_count == 2 - expected_request = _make_begin_request(project, database=database) + expected_request = _make_begin_request(project, database=database_id) begin_txn.assert_called_with(request=expected_request) commit_method = ds_api.commit assert commit_method.call_count == 2 mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL - commit_method.assert_called_with( - request={ - "project_id": project, - "database_id": database, - "mode": mode, - "mutations": [], - "transaction": id_, - } - ) + expected_request = { + "project_id": project, + "mode": mode, + "mutations": [], + "transaction": id_, + } + _set_database_id_to_request(expected_request, database_id) + + commit_method.assert_called_with(request=expected_request) ds_api.rollback.assert_not_called() -def test_transaction_begin(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_begin(database_id): project = "PROJECT" - database = "DATABASE" id_ = 889 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, database=database, datastore_api=ds_api) + client = _Client(project, database=database_id, datastore_api=ds_api) xact = _make_transaction(client) xact.begin() assert xact.id == id_ - expected_request = _make_begin_request(project, database=database) + expected_request = _make_begin_request(project, database=database_id) + ds_api.begin_transaction.assert_called_once_with(request=expected_request) -def test_transaction_begin_w_readonly(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_begin_w_readonly(database_id): project = "PROJECT" id_ = 889 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client, read_only=True) xact.begin() assert xact.id == id_ - expected_request = _make_begin_request(project, read_only=True) + expected_request = _make_begin_request( + project, read_only=True, database=database_id + ) ds_api.begin_transaction.assert_called_once_with(request=expected_request) -def test_transaction_begin_w_read_time(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_begin_w_read_time(database_id): from datetime import datetime project = "PROJECT" id_ = 889 read_time = datetime.utcfromtimestamp(1641058200.123456) ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client, read_only=True, read_time=read_time) xact.begin() assert xact.id == id_ - expected_request = _make_begin_request(project, read_only=True, read_time=read_time) + expected_request = _make_begin_request( + project, read_only=True, read_time=read_time, database=database_id + ) ds_api.begin_transaction.assert_called_once_with(request=expected_request) -def test_transaction_begin_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_begin_w_retry_w_timeout(database_id): project = "PROJECT" id_ = 889 retry = mock.Mock() timeout = 100000 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) xact.begin(retry=retry, timeout=timeout) assert xact.id == id_ - expected_request = _make_begin_request(project) + expected_request = _make_begin_request(project, database=database_id) ds_api.begin_transaction.assert_called_once_with( request=expected_request, retry=retry, @@ -201,37 +216,38 @@ def test_transaction_begin_w_retry_w_timeout(): ) -def test_transaction_begin_tombstoned(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_begin_tombstoned(database_id): project = "PROJECT" id_ = 1094 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) xact.begin() assert xact.id == id_ - expected_request = _make_begin_request(project) + expected_request = _make_begin_request(project, database=database_id) ds_api.begin_transaction.assert_called_once_with(request=expected_request) xact.rollback() - - client._datastore_api.rollback.assert_called_once_with( - request={"project_id": project, "database_id": "", "transaction": id_} - ) + expected_request = {"project_id": project, "transaction": id_} + _set_database_id_to_request(expected_request, database_id) + client._datastore_api.rollback.assert_called_once_with(request=expected_request) assert xact.id is None with pytest.raises(ValueError): xact.begin() -def test_transaction_begin_w_begin_transaction_failure(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_begin_w_begin_transaction_failure(database_id): project = "PROJECT" id_ = 712 ds_api = _make_datastore_api(xact_id=id_) ds_api.begin_transaction = mock.Mock(side_effect=RuntimeError, spec=[]) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) with pytest.raises(RuntimeError): @@ -239,103 +255,108 @@ def test_transaction_begin_w_begin_transaction_failure(): assert xact.id is None - expected_request = _make_begin_request(project) + expected_request = _make_begin_request(project, database=database_id) ds_api.begin_transaction.assert_called_once_with(request=expected_request) -def test_transaction_rollback(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_rollback(database_id): project = "PROJECT" - database = "DATABASE" id_ = 239 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, database=database, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) xact.begin() xact.rollback() assert xact.id is None - ds_api.rollback.assert_called_once_with( - request={"project_id": project, "database_id": database, "transaction": id_} - ) + expected_request = {"project_id": project, "transaction": id_} + _set_database_id_to_request(expected_request, database_id) + ds_api.rollback.assert_called_once_with(request=expected_request) -def test_transaction_rollback_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_rollback_w_retry_w_timeout(database_id): project = "PROJECT" id_ = 239 retry = mock.Mock() timeout = 100000 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) xact.begin() xact.rollback(retry=retry, timeout=timeout) assert xact.id is None + expected_request = {"project_id": project, "transaction": id_} + _set_database_id_to_request(expected_request, database_id) + ds_api.rollback.assert_called_once_with( - request={"project_id": project, "database_id": "", "transaction": id_}, + request=expected_request, retry=retry, timeout=timeout, ) -def test_transaction_commit_no_partial_keys(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_commit_no_partial_keys(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" - database = "DATABASE" id_ = 1002930 mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, database=database, datastore_api=ds_api) + client = _Client(project, database=database_id, datastore_api=ds_api) xact = _make_transaction(client) xact.begin() xact.commit() - ds_api.commit.assert_called_once_with( - request={ - "project_id": project, - "database_id": database, - "mode": mode, - "mutations": [], - "transaction": id_, - } - ) + expected_request = { + "project_id": project, + "mode": mode, + "mutations": [], + "transaction": id_, + } + _set_database_id_to_request(expected_request, database_id) + ds_api.commit.assert_called_once_with(request=expected_request) assert xact.id is None -def test_transaction_commit_w_partial_keys_w_retry_w_timeout(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_commit_w_partial_keys_w_retry_w_timeout(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" kind = "KIND" id1 = 123 mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL - key = _make_key(kind, id1, project) + key = _make_key(kind, id1, project, database=database_id) id2 = 234 retry = mock.Mock() timeout = 100000 ds_api = _make_datastore_api(key, xact_id=id2) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) xact.begin() - entity = _Entity() + entity = _Entity(database=database_id) xact.put(entity) xact.commit(retry=retry, timeout=timeout) + expected_request = { + "project_id": project, + "mode": mode, + "mutations": xact.mutations, + "transaction": id2, + } + _set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": xact.mutations, - "transaction": id2, - }, + request=expected_request, retry=retry, timeout=timeout, ) @@ -343,13 +364,14 @@ def test_transaction_commit_w_partial_keys_w_retry_w_timeout(): assert entity.key.path == [{"kind": kind, "id": id1}] -def test_transaction_context_manager_no_raise(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_context_manager_no_raise(database_id): from google.cloud.datastore_v1.types import datastore as datastore_pb2 project = "PROJECT" id_ = 912830 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) with xact: @@ -358,29 +380,32 @@ def test_transaction_context_manager_no_raise(): assert xact.id is None - expected_request = _make_begin_request(project) + expected_request = _make_begin_request(project, database=database_id) ds_api.begin_transaction.assert_called_once_with(request=expected_request) mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL + expected_request = { + "project_id": project, + "mode": mode, + "mutations": [], + "transaction": id_, + } + _set_database_id_to_request(expected_request, database_id) + client._datastore_api.commit.assert_called_once_with( - request={ - "project_id": project, - "database_id": "", - "mode": mode, - "mutations": [], - "transaction": id_, - }, + request=expected_request, ) -def test_transaction_context_manager_w_raise(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_context_manager_w_raise(database_id): class Foo(Exception): pass project = "PROJECT" id_ = 614416 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) + client = _Client(project, datastore_api=ds_api, database=database_id) xact = _make_transaction(client) xact._mutation = object() try: @@ -392,22 +417,23 @@ class Foo(Exception): assert xact.id is None - expected_request = _make_begin_request(project) + expected_request = _make_begin_request(project, database=database_id) + _set_database_id_to_request(expected_request, database_id) ds_api.begin_transaction.assert_called_once_with(request=expected_request) client._datastore_api.commit.assert_not_called() - - client._datastore_api.rollback.assert_called_once_with( - request={"project_id": project, "database_id": "", "transaction": id_} - ) + expected_request = {"project_id": project, "transaction": id_} + _set_database_id_to_request(expected_request, database_id) + client._datastore_api.rollback.assert_called_once_with(request=expected_request) -def test_transaction_put_read_only(): +@pytest.mark.parametrize("database_id", [None, "somedb"]) +def test_transaction_put_read_only(database_id): project = "PROJECT" id_ = 943243 ds_api = _make_datastore_api(xact_id=id_) - client = _Client(project, datastore_api=ds_api) - entity = _Entity() + client = _Client(project, datastore_api=ds_api, database=database_id) + entity = _Entity(database=database_id) xact = _make_transaction(client, read_only=True) xact.begin() @@ -415,7 +441,7 @@ def test_transaction_put_read_only(): xact.put(entity) -def _make_key(kind, id_, project, database=""): +def _make_key(kind, id_, project, database=None): from google.cloud.datastore_v1.types import entity as entity_pb2 key = entity_pb2.Key() @@ -428,15 +454,15 @@ def _make_key(kind, id_, project, database=""): class _Entity(dict): - def __init__(self): + def __init__(self, database=None): super(_Entity, self).__init__() from google.cloud.datastore.key import Key - self.key = Key("KIND", project="PROJECT") + self.key = Key("KIND", project="PROJECT", database=database) class _Client(object): - def __init__(self, project, datastore_api=None, namespace=None, database=""): + def __init__(self, project, datastore_api=None, namespace=None, database=None): self.project = project if datastore_api is None: datastore_api = _make_datastore_api() @@ -495,13 +521,14 @@ def _make_transaction(client, **kw): return Transaction(client, **kw) -def _make_begin_request(project, read_only=False, read_time=None, database=""): +def _make_begin_request(project, read_only=False, read_time=None, database=None): expected_options = _make_options(read_only=read_only, read_time=read_time) - return { + request = { "project_id": project, - "database_id": database, "transaction_options": expected_options, } + _set_database_id_to_request(request, database) + return request def _make_commit_response(*keys): From 1b7f92c6a3bc57dc49572b78679d922e6a1a04d4 Mon Sep 17 00:00:00 2001 From: Mariatta Wijaya Date: Mon, 8 May 2023 22:35:27 -0700 Subject: [PATCH 2/8] Add test case for when parent database name != child database name --- tests/unit/test_key.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_key.py b/tests/unit/test_key.py index e78981f2..517013d5 100644 --- a/tests/unit/test_key.py +++ b/tests/unit/test_key.py @@ -114,7 +114,7 @@ def test_key_ctor_parent_bad_namespace(database_id): parent_key = _make_key( "KIND", 1234, namespace="FOO", project=_DEFAULT_PROJECT, database=database_id ) - with pytest.raises(ValueError): + with pytest.raises(ValueError) as exc: _make_key( "KIND2", 1234, @@ -123,13 +123,28 @@ def test_key_ctor_parent_bad_namespace(database_id): PROJECT=_DEFAULT_PROJECT, database=database_id, ) + assert "Child namespace must agree with parent's." in str(exc.value) @pytest.mark.parametrize("database_id", [None, "somedb"]) def test_key_ctor_parent_bad_project(database_id): parent_key = _make_key("KIND", 1234, project="FOO", database=database_id) - with pytest.raises(ValueError): + with pytest.raises(ValueError) as exc: _make_key("KIND2", 1234, parent=parent_key, project="BAR", database=database_id) + assert "Child project must agree with parent's." in str(exc.value) + + +def test_key_ctor_parent_bad_database(): + parent_key = _make_key("KIND", 1234, project=_DEFAULT_PROJECT, database="db1") + with pytest.raises(ValueError) as exc: + _make_key( + "KIND2", + 1234, + parent=parent_key, + PROJECT=_DEFAULT_PROJECT, + database="db2", + ) + assert "Child database must agree with parent's" in str(exc.value) def test_key_ctor_parent_empty_path(): From 544dab720d9c39c7c6049c7372f54815d7df8be9 Mon Sep 17 00:00:00 2001 From: Mariatta Wijaya Date: Mon, 8 May 2023 22:39:12 -0700 Subject: [PATCH 3/8] Update owlbot, removing the named db parameter --- owlbot.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/owlbot.py b/owlbot.py index 8186f515..d5040231 100644 --- a/owlbot.py +++ b/owlbot.py @@ -120,8 +120,7 @@ def system\(session\): """\ @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) @nox.parametrize("disable_grpc", [False, True]) -@nox.parametrize("use_named_db", [False, True]) -def system(session, disable_grpc, use_named_db): +def system(session, disable_grpc): """, ) @@ -134,7 +133,6 @@ def system(session, disable_grpc, use_named_db): env = {} if disable_grpc: env["GOOGLE_CLOUD_DISABLE_GRPC"] = "True" - env["SYSTEM_TESTS_DATABASE"] = "system-tests-named-db" if use_named_db else "" # Run py.test against the system tests. """, From da7d504cba004aabb0a7945e6d11f3906fc0f4de Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Tue, 13 Jun 2023 10:40:00 +0000 Subject: [PATCH 4/8] Reverted test fixes --- google/cloud/datastore/_http.py | 69 +++++++++++++++++++-------------- tests/unit/test__http.py | 6 ++- 2 files changed, 44 insertions(+), 31 deletions(-) diff --git a/google/cloud/datastore/_http.py b/google/cloud/datastore/_http.py index 4090a475..fbbecab7 100644 --- a/google/cloud/datastore/_http.py +++ b/google/cloud/datastore/_http.py @@ -52,15 +52,6 @@ def _make_request_pb(request, request_pb_type): return request -def _make_routing_header(request): - """Helper for creating a routing header from the request body""" - database_id = request.database_id if "database_id" in request else "" - header = { - "x-goog-request-params": f"project_id={request.project_id}&database_id={database_id}" - } - return header - - def _request( http, project, @@ -68,7 +59,7 @@ def _request( data, base_url, client_info, - headers=None, + database=None, retry=None, timeout=None, ): @@ -94,8 +85,8 @@ def _request( :type client_info: :class:`google.api_core.client_info.ClientInfo` :param client_info: used to generate user agent. - :type headers: :class:`dict` - :param headers: (Optional) custom headers for the request. + :type database: str + :param database: (Optional) The database to make the request for. :type retry: :class:`google.api_core.retry.Retry` :param retry: (Optional) retry policy for the request @@ -109,14 +100,12 @@ def _request( response code is not 200 OK. """ user_agent = client_info.to_user_agent() - req_headers = { + headers = { "Content-Type": "application/x-protobuf", "User-Agent": user_agent, connection_module.CLIENT_INFO_HEADER: user_agent, } - if headers: - req_headers.update(headers) - + _update_headers(headers, project, database) api_url = build_api_url(project, method, base_url) requester = http.request @@ -128,12 +117,12 @@ def _request( response = requester( url=api_url, method="POST", - headers=req_headers, + headers=headers, data=data, timeout=timeout, ) else: - response = requester(url=api_url, method="POST", headers=req_headers, data=data) + response = requester(url=api_url, method="POST", headers=headers, data=data) if response.status_code != 200: error_status = status_pb2.Status.FromString(response.content) @@ -152,7 +141,7 @@ def _rpc( client_info, request_pb, response_pb_cls, - headers=None, + database=None, retry=None, timeout=None, ): @@ -182,8 +171,8 @@ def _rpc( :param response_pb_cls: The class used to unmarshall the response protobuf. - :type headers: :class:`dict` - :param headers: (Optional) custom headers for the request. + :type database: str + :param database: (Optional) The database to make the request for. :type retry: :class:`google.api_core.retry.Retry` :param retry: (Optional) retry policy for the request @@ -197,7 +186,7 @@ def _rpc( req_data = request_pb._pb.SerializeToString() kwargs = _make_retry_timeout_kwargs(retry, timeout) response = _request( - http, project, method, req_data, base_url, client_info, headers, **kwargs + http, project, method, req_data, base_url, client_info, database, **kwargs ) return response_pb_cls.deserialize(response) @@ -256,6 +245,7 @@ def lookup(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.LookupRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -265,7 +255,7 @@ def lookup(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.LookupResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -288,6 +278,7 @@ def run_query(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.RunQueryRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -297,7 +288,7 @@ def run_query(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.RunQueryResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -322,6 +313,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None): request, _datastore_pb2.RunAggregationQueryRequest ) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -331,7 +323,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.RunAggregationQueryResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -354,6 +346,7 @@ def begin_transaction(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.BeginTransactionRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -363,7 +356,7 @@ def begin_transaction(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.BeginTransactionResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -386,6 +379,7 @@ def commit(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.CommitRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -395,7 +389,7 @@ def commit(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.CommitResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -418,6 +412,7 @@ def rollback(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.RollbackRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -427,7 +422,7 @@ def rollback(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.RollbackResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -450,6 +445,7 @@ def allocate_ids(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.AllocateIdsRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -459,7 +455,7 @@ def allocate_ids(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.AllocateIdsResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) @@ -482,6 +478,7 @@ def reserve_ids(self, request, retry=None, timeout=None): """ request_pb = _make_request_pb(request, _datastore_pb2.ReserveIdsRequest) project_id = request_pb.project_id + database_id = request_pb.database_id return _rpc( self.client._http, @@ -491,7 +488,19 @@ def reserve_ids(self, request, retry=None, timeout=None): self.client._client_info, request_pb, _datastore_pb2.ReserveIdsResponse, - _make_routing_header(request_pb), + database_id, retry=retry, timeout=timeout, ) + + +def _update_headers(headers, project_id, database_id=None): + """Update the request headers. + Pass the project id, or optionally the database_id if provided. + """ + if database_id: + headers[ + "x-goog-request-params" + ] = f"project_id={project_id}&database_id={database_id}" + else: + headers["x-goog-request-params"] = f"project_id={project_id}" diff --git a/tests/unit/test__http.py b/tests/unit/test__http.py index 8a3a8df6..6f1affd6 100644 --- a/tests/unit/test__http.py +++ b/tests/unit/test__http.py @@ -122,6 +122,7 @@ def _request_helper(retry=None, timeout=None): "Content-Type": "application/x-protobuf", "User-Agent": user_agent, connection_module.CLIENT_INFO_HEADER: user_agent, + "x-goog-request-params": f"project_id={project}", } if retry is not None: @@ -1026,11 +1027,14 @@ def _make_client_info(user_agent=_USER_AGENT): def _verify_protobuf_call(http, expected_url, pb, retry=None, timeout=None): from google.cloud import _http as connection_module + routing_header = f"project_id={pb.project_id}" + if pb.database_id: + routing_header += f"&database_id={pb.database_id}" expected_headers = { "Content-Type": "application/x-protobuf", "User-Agent": _USER_AGENT, connection_module.CLIENT_INFO_HEADER: _USER_AGENT, - "x-goog-request-params": f"project_id={pb.project_id}&database_id={pb.database_id}", + "x-goog-request-params": routing_header, } if retry is not None: From cbb59fb5f86e5c38874461f48f1b1b47e27097dd Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Tue, 13 Jun 2023 10:58:07 +0000 Subject: [PATCH 5/8] fixing tests --- google/cloud/datastore/_http.py | 1 - tests/unit/test__http.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/google/cloud/datastore/_http.py b/google/cloud/datastore/_http.py index fc58f22b..3271a5cf 100644 --- a/google/cloud/datastore/_http.py +++ b/google/cloud/datastore/_http.py @@ -247,7 +247,6 @@ def lookup(self, request, retry=None, timeout=None): project_id = request_pb.project_id database_id = request_pb.database_id - database = request_pb.database_id return _rpc( self.client._http, project_id, diff --git a/tests/unit/test__http.py b/tests/unit/test__http.py index 96ddfc0d..17f19f77 100644 --- a/tests/unit/test__http.py +++ b/tests/unit/test__http.py @@ -177,7 +177,7 @@ def test__request_failure(): ) with pytest.raises(BadRequest) as exc: - _request(session, project, method, data, uri, client_info) + _request(session, project, method, data, uri, client_info, None) expected_message = "400 Entity value is indexed." assert exc.match(expected_message) From b09aae8a66ceefad6a813e35f20bb3bd5c2e776a Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Tue, 13 Jun 2023 11:47:33 +0000 Subject: [PATCH 6/8] fix code coverage --- tests/unit/test__http.py | 25 +------------------------ 1 file changed, 1 insertion(+), 24 deletions(-) diff --git a/tests/unit/test__http.py b/tests/unit/test__http.py index 148e6d20..4c8f9eca 100644 --- a/tests/unit/test__http.py +++ b/tests/unit/test__http.py @@ -1047,28 +1047,6 @@ def test_update_headers_with_database_id(): ) -def test_update_headers_without_database_id(): - from google.cloud.datastore._http import _update_headers - - headers = {} - project_id = "someproject" - _update_headers(headers, project_id) - assert headers["x-goog-request-params"] == f"project_id={project_id}" - - -def test_update_headers_with_database_id(): - from google.cloud.datastore._http import _update_headers - - headers = {} - project_id = "someproject" - database_id = "somedb" - _update_headers(headers, project_id, database_id=database_id) - assert ( - headers["x-goog-request-params"] - == f"project_id={project_id}&database_id={database_id}" - ) - - def _make_http_datastore_api(*args, **kwargs): from google.cloud.datastore._http import HTTPDatastoreAPI @@ -1122,12 +1100,11 @@ def _verify_protobuf_call( from google.cloud import _http as connection_module from google.cloud.datastore._http import _update_headers - routing_header = f"project_id={pb.project_id}" expected_headers = { "Content-Type": "application/x-protobuf", "User-Agent": _USER_AGENT, connection_module.CLIENT_INFO_HEADER: _USER_AGENT, - "x-goog-request-params": routing_header, + "x-goog-request-params": f"project_id={pb.project_id}", } _update_headers(expected_headers, project, database_id=database) if retry is not None: From 2be39dab9b154623a304cb16e9ed8367b9e4c197 Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Thu, 15 Jun 2023 09:08:01 +0000 Subject: [PATCH 7/8] pr suggestion --- google/cloud/datastore/key.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/datastore/key.py b/google/cloud/datastore/key.py index b0c8e5d6..4384131c 100644 --- a/google/cloud/datastore/key.py +++ b/google/cloud/datastore/key.py @@ -171,7 +171,7 @@ def __hash__(self): """ hash_val = hash(self.flat_path) + hash(self.project) + hash(self.namespace) if self.database: - hash_val + hash(self.database) + hash_val = hash_val + hash(self.database) return hash_val @staticmethod From a7fbbdab4897f4a72cb991d91d4f87a8875e7eb9 Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Tue, 20 Jun 2023 18:49:37 +0000 Subject: [PATCH 8/8] address pr comments --- google/cloud/datastore/aggregation.py | 4 +-- google/cloud/datastore/batch.py | 2 +- google/cloud/datastore/client.py | 8 +++--- google/cloud/datastore/helpers.py | 2 +- google/cloud/datastore/query.py | 4 +-- google/cloud/datastore/transaction.py | 6 ++-- tests/unit/test__http.py | 14 +++++----- tests/unit/test_aggregation.py | 4 +-- tests/unit/test_batch.py | 14 +++++----- tests/unit/test_client.py | 40 +++++++++++++-------------- tests/unit/test_query.py | 6 ++-- tests/unit/test_transaction.py | 22 +++++++-------- 12 files changed, 63 insertions(+), 63 deletions(-) diff --git a/google/cloud/datastore/aggregation.py b/google/cloud/datastore/aggregation.py index 25fbf992..421ffc93 100644 --- a/google/cloud/datastore/aggregation.py +++ b/google/cloud/datastore/aggregation.py @@ -393,7 +393,7 @@ def _next_page(self): "read_options": read_options, "aggregation_query": query_pb, } - helpers._set_database_id_to_request(request, self.client.database) + helpers.set_database_id_to_request(request, self.client.database) response_pb = self.client._datastore_api.run_aggregation_query( request=request, **kwargs, @@ -414,7 +414,7 @@ def _next_page(self): "read_options": read_options, "aggregation_query": query_pb, } - helpers._set_database_id_to_request(request, self.client.database) + helpers.set_database_id_to_request(request, self.client.database) response_pb = self.client._datastore_api.run_aggregation_query( request=request, **kwargs, diff --git a/google/cloud/datastore/batch.py b/google/cloud/datastore/batch.py index 276c0822..e0dbf26d 100644 --- a/google/cloud/datastore/batch.py +++ b/google/cloud/datastore/batch.py @@ -303,7 +303,7 @@ def _commit(self, retry, timeout): "mutations": self._mutations, } - helpers._set_database_id_to_request(request, self._client.database) + helpers.set_database_id_to_request(request, self._client.database) commit_response_pb = self._client._datastore_api.commit( request=request, diff --git a/google/cloud/datastore/client.py b/google/cloud/datastore/client.py index 8b3d89d4..fe25a0e0 100644 --- a/google/cloud/datastore/client.py +++ b/google/cloud/datastore/client.py @@ -208,7 +208,7 @@ def _extended_lookup( "keys": key_pbs, "read_options": read_options, } - helpers._set_database_id_to_request(request, database) + helpers.set_database_id_to_request(request, database) lookup_response = datastore_api.lookup( request=request, **kwargs, @@ -761,7 +761,7 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None): "project_id": incomplete_key.project, "keys": incomplete_key_pbs, } - helpers._set_database_id_to_request(request, self.database) + helpers.set_database_id_to_request(request, self.database) response_pb = self._datastore_api.allocate_ids( request=request, **kwargs, @@ -1006,7 +1006,7 @@ def reserve_ids_sequential(self, complete_key, num_ids, retry=None, timeout=None "project_id": complete_key.project, "keys": key_pbs, } - helpers._set_database_id_to_request(request, self.database) + helpers.set_database_id_to_request(request, self.database) self._datastore_api.reserve_ids( request=request, **kwargs, @@ -1059,7 +1059,7 @@ def reserve_ids_multi(self, complete_keys, retry=None, timeout=None): "project_id": complete_keys[0].project, "keys": key_pbs, } - helpers._set_database_id_to_request(request, complete_keys[0].database) + helpers.set_database_id_to_request(request, complete_keys[0].database) self._datastore_api.reserve_ids( request=request, diff --git a/google/cloud/datastore/helpers.py b/google/cloud/datastore/helpers.py index fa6730ce..2deecabe 100644 --- a/google/cloud/datastore/helpers.py +++ b/google/cloud/datastore/helpers.py @@ -490,7 +490,7 @@ def _set_protobuf_value(value_pb, val): setattr(value_pb, attr, val) -def _set_database_id_to_request(request, database_id=None): +def set_database_id_to_request(request, database_id=None): """ Set the "database_id" field to the request only if it was provided. """ diff --git a/google/cloud/datastore/query.py b/google/cloud/datastore/query.py index 14335268..289605bb 100644 --- a/google/cloud/datastore/query.py +++ b/google/cloud/datastore/query.py @@ -809,7 +809,7 @@ def _next_page(self): "query": query_pb, } - helpers._set_database_id_to_request(request, self.client.database) + helpers.set_database_id_to_request(request, self.client.database) response_pb = self.client._datastore_api.run_query( request=request, @@ -836,7 +836,7 @@ def _next_page(self): "read_options": read_options, "query": query_pb, } - helpers._set_database_id_to_request(request, self.client.database) + helpers.set_database_id_to_request(request, self.client.database) response_pb = self.client._datastore_api.run_query( request=request, diff --git a/google/cloud/datastore/transaction.py b/google/cloud/datastore/transaction.py index 73a9c68e..3e71ae26 100644 --- a/google/cloud/datastore/transaction.py +++ b/google/cloud/datastore/transaction.py @@ -18,7 +18,7 @@ from google.cloud.datastore_v1.types import TransactionOptions from google.protobuf import timestamp_pb2 -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request def _make_retry_timeout_kwargs(retry, timeout): @@ -229,7 +229,7 @@ def begin(self, retry=None, timeout=None): "project_id": self.project, "transaction_options": self._options, } - _set_database_id_to_request(request, self._client.database) + set_database_id_to_request(request, self._client.database) try: response_pb = self._client._datastore_api.begin_transaction( @@ -267,7 +267,7 @@ def rollback(self, retry=None, timeout=None): "transaction": self._id, } - _set_database_id_to_request(request, self._client.database) + set_database_id_to_request(request, self._client.database) self._client._datastore_api.rollback(request=request, **kwargs) finally: super(Transaction, self).rollback() diff --git a/tests/unit/test__http.py b/tests/unit/test__http.py index 4c8f9eca..48e7f5b6 100644 --- a/tests/unit/test__http.py +++ b/tests/unit/test__http.py @@ -18,7 +18,7 @@ import pytest import requests -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request def test__make_retry_timeout_kwargs_w_empty(): @@ -297,7 +297,7 @@ def _lookup_single_helper( "keys": [key_pb], "read_options": read_options, } - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.lookup(request=request, **kwargs) @@ -690,7 +690,7 @@ def _begin_transaction_helper(options=None, retry=None, timeout=None, database=N # Make request. ds_api = _make_http_datastore_api(client) request = {"project_id": project} - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) if options is not None: request["transaction_options"] = options @@ -768,7 +768,7 @@ def _commit_helper(transaction=None, retry=None, timeout=None, database=None): ds_api = _make_http_datastore_api(client) request = {"project_id": project, "mutations": [mutation]} - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) if transaction is not None: request["transaction"] = transaction mode = request["mode"] = rq_class.Mode.TRANSACTIONAL @@ -846,7 +846,7 @@ def _rollback_helper(retry=None, timeout=None, database=None): # Make request. ds_api = _make_http_datastore_api(client) request = {"project_id": project, "transaction": transaction} - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.rollback(request=request, **kwargs) @@ -912,7 +912,7 @@ def _allocate_ids_helper(count=0, retry=None, timeout=None, database=None): ds_api = _make_http_datastore_api(client) request = {"project_id": project, "keys": before_key_pbs} - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.allocate_ids(request=request, **kwargs) @@ -981,7 +981,7 @@ def _reserve_ids_helper(count=0, retry=None, timeout=None, database=None): ds_api = _make_http_datastore_api(client) request = {"project_id": project, "keys": before_key_pbs} - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) kwargs = _retry_timeout_kw(retry, timeout, http) response = ds_api.reserve_ids(request=request, **kwargs) diff --git a/tests/unit/test_aggregation.py b/tests/unit/test_aggregation.py index a45b05b2..ebfa9a3f 100644 --- a/tests/unit/test_aggregation.py +++ b/tests/unit/test_aggregation.py @@ -16,7 +16,7 @@ import pytest from google.cloud.datastore.aggregation import CountAggregation, AggregationQuery -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request from tests.unit.test_query import _make_query, _make_client @@ -398,7 +398,7 @@ def _next_page_helper(txn_id=None, retry=None, timeout=None, database_id=None): "read_options": read_options, "aggregation_query": aggregation_query._to_pb(), } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) expected_call = mock.call( request=expected_request, **kwargs, diff --git a/tests/unit/test_batch.py b/tests/unit/test_batch.py index 25a968fc..67f5cff5 100644 --- a/tests/unit/test_batch.py +++ b/tests/unit/test_batch.py @@ -18,7 +18,7 @@ import mock import pytest -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request def _make_batch(client): @@ -79,7 +79,7 @@ def test_batch_current(database_id): "mutations": [], "transaction": None, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) commit_method.assert_called_with(request=expected_request) @@ -334,7 +334,7 @@ def _batch_commit_helper(timeout=None, retry=None, database=None): "mutations": [], "transaction": None, } - _set_database_id_to_request(expected_request, database) + set_database_id_to_request(expected_request, database) commit_method.assert_called_with(request=expected_request, **kwargs) @@ -383,7 +383,7 @@ def test_batch_commit_w_partial_key_entity(database_id): "mutations": [], "transaction": None, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with(request=expected_request) assert not entity.key.is_partial @@ -419,7 +419,7 @@ def test_batch_as_context_mgr_wo_error(database_id): "mutations": batch.mutations, "transaction": None, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) commit_method.assert_called_with(request=expected_request) @@ -471,8 +471,8 @@ def test_batch_as_context_mgr_nested(database_id): "mutations": batch1.mutations, "transaction": None, } - _set_database_id_to_request(expected_request_1, database_id) - _set_database_id_to_request(expected_request_2, database_id) + set_database_id_to_request(expected_request_1, database_id) + set_database_id_to_request(expected_request_2, database_id) commit_method.assert_called_with(request=expected_request_1) commit_method.assert_called_with(request=expected_request_2) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 06cfa37f..119bab79 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -18,7 +18,7 @@ import mock import pytest -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request PROJECT = "dummy-project-123" DATABASE = "dummy-database-123" @@ -451,7 +451,7 @@ def test_client_get_multi_miss(database_id): "keys": [key.to_protobuf()], "read_options": read_options, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.lookup.assert_called_once_with(request=expected_request) @@ -492,7 +492,7 @@ def test_client_get_multi_miss_w_missing(database_id): "keys": [key_pb], "read_options": read_options, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.lookup.assert_called_once_with(request=expected_request) @@ -545,7 +545,7 @@ def test_client_get_multi_miss_w_deferred(database_id): "keys": [key_pb], "read_options": read_options, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.lookup.assert_called_once_with(request=expected_request) @@ -601,7 +601,7 @@ def test_client_get_multi_w_deferred_from_backend_but_not_passed(database_id): "keys": [key2_pb], "read_options": read_options, } - _set_database_id_to_request(expected_request_1, database_id) + set_database_id_to_request(expected_request_1, database_id) ds_api.lookup.assert_any_call( request=expected_request_1, ) @@ -611,7 +611,7 @@ def test_client_get_multi_w_deferred_from_backend_but_not_passed(database_id): "keys": [key1_pb, key2_pb], "read_options": read_options, } - _set_database_id_to_request(expected_request_2, database_id) + set_database_id_to_request(expected_request_2, database_id) ds_api.lookup.assert_any_call( request=expected_request_2, ) @@ -702,7 +702,7 @@ def test_client_get_multi_hit_w_transaction(database_id): "keys": [key.to_protobuf()], "read_options": read_options, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.lookup.assert_called_once_with(request=expected_request) @@ -747,7 +747,7 @@ def test_client_get_multi_hit_w_read_time(database_id): "keys": [key.to_protobuf()], "read_options": read_options, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.lookup.assert_called_once_with(request=expected_request) @@ -787,7 +787,7 @@ def test_client_get_multi_hit_multiple_keys_same_project(database_id): "keys": [key1.to_protobuf(), key2.to_protobuf()], "read_options": read_options, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.lookup.assert_called_once_with(request=expected_request) @@ -911,7 +911,7 @@ def test_client_put_multi_no_batch_w_partial_key_w_retry_w_timeout(database_id): "mutations": mock.ANY, "transaction": None, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with( request=expected_request, retry=retry, @@ -1007,7 +1007,7 @@ def test_client_delete_multi_no_batch_w_retry_w_timeout(database_id): "mutations": mock.ANY, "transaction": None, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with( request=expected_request, retry=retry, @@ -1103,7 +1103,7 @@ def test_client_allocate_ids_w_partial_key(database_id): expected_keys = [incomplete_key.to_protobuf()] * num_ids expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) alloc_ids.assert_called_once_with(request=expected_request) @@ -1129,7 +1129,7 @@ def test_client_allocate_ids_w_partial_key_w_retry_w_timeout(database_id): expected_keys = [incomplete_key.to_protobuf()] * num_ids expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) alloc_ids.assert_called_once_with( request=expected_request, retry=retry, @@ -1155,7 +1155,7 @@ def test_client_reserve_ids_sequential_w_completed_key(database_id): ) expected_keys = [key.to_protobuf() for key in reserved_keys] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with(request=expected_request) @@ -1180,7 +1180,7 @@ def test_client_reserve_ids_sequential_w_completed_key_w_retry_w_timeout(databas ) expected_keys = [key.to_protobuf() for key in reserved_keys] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with( request=expected_request, retry=retry, @@ -1207,7 +1207,7 @@ def test_client_reserve_ids_sequential_w_completed_key_w_ancestor(database_id): ) expected_keys = [key.to_protobuf() for key in reserved_keys] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with(request=expected_request) @@ -1313,7 +1313,7 @@ def test_client_reserve_ids_w_completed_key(database_id): ) expected_keys = [key.to_protobuf() for key in reserved_keys] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with(request=expected_request) _assert_reserve_ids_warning(warned) @@ -1343,7 +1343,7 @@ def test_client_reserve_ids_w_completed_key_w_retry_w_timeout(database_id): ) expected_keys = [key.to_protobuf() for key in reserved_keys] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with( request=expected_request, retry=retry, @@ -1374,7 +1374,7 @@ def test_client_reserve_ids_w_completed_key_w_ancestor(database_id): ) expected_keys = [key.to_protobuf() for key in reserved_keys] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with(request=expected_request) @@ -1706,7 +1706,7 @@ def test_client_reserve_ids_multi(database_id): expected_keys = [key1.to_protobuf(), key2.to_protobuf()] expected_request = {"project_id": PROJECT, "keys": expected_keys} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) reserve_ids.assert_called_once_with(request=expected_request) diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py index 2b84c731..25b3febb 100644 --- a/tests/unit/test_query.py +++ b/tests/unit/test_query.py @@ -25,7 +25,7 @@ BaseCompositeFilter, ) -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request _PROJECT = "PROJECT" @@ -880,7 +880,7 @@ def _next_page_helper( "read_options": read_options, "query": empty_query, } - _set_database_id_to_request(expected_request, database) + set_database_id_to_request(expected_request, database) ds_api.run_query.assert_called_once_with( request=expected_request, **kwargs, @@ -977,7 +977,7 @@ def test_iterator__next_page_w_skipped_lt_offset(database_id): "read_options": read_options, "query": query, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) expected_calls.append(mock.call(request=expected_request)) assert ds_api.run_query.call_args_list == expected_calls diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py index 96854db6..23574ef4 100644 --- a/tests/unit/test_transaction.py +++ b/tests/unit/test_transaction.py @@ -15,7 +15,7 @@ import mock import pytest -from google.cloud.datastore.helpers import _set_database_id_to_request +from google.cloud.datastore.helpers import set_database_id_to_request @pytest.mark.parametrize("database_id", [None, "somedb"]) @@ -130,7 +130,7 @@ def test_transaction_current(database_id): "mutations": [], "transaction": id_, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) commit_method.assert_called_with(request=expected_request) @@ -233,7 +233,7 @@ def test_transaction_begin_tombstoned(database_id): xact.rollback() expected_request = {"project_id": project, "transaction": id_} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) client._datastore_api.rollback.assert_called_once_with(request=expected_request) assert xact.id is None @@ -272,7 +272,7 @@ def test_transaction_rollback(database_id): assert xact.id is None expected_request = {"project_id": project, "transaction": id_} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.rollback.assert_called_once_with(request=expected_request) @@ -292,7 +292,7 @@ def test_transaction_rollback_w_retry_w_timeout(database_id): assert xact.id is None expected_request = {"project_id": project, "transaction": id_} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.rollback.assert_called_once_with( request=expected_request, @@ -321,7 +321,7 @@ def test_transaction_commit_no_partial_keys(database_id): "mutations": [], "transaction": id_, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with(request=expected_request) assert xact.id is None @@ -353,7 +353,7 @@ def test_transaction_commit_w_partial_keys_w_retry_w_timeout(database_id): "mutations": xact.mutations, "transaction": id2, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.commit.assert_called_once_with( request=expected_request, @@ -390,7 +390,7 @@ def test_transaction_context_manager_no_raise(database_id): "mutations": [], "transaction": id_, } - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) client._datastore_api.commit.assert_called_once_with( request=expected_request, @@ -418,12 +418,12 @@ class Foo(Exception): assert xact.id is None expected_request = _make_begin_request(project, database=database_id) - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) ds_api.begin_transaction.assert_called_once_with(request=expected_request) client._datastore_api.commit.assert_not_called() expected_request = {"project_id": project, "transaction": id_} - _set_database_id_to_request(expected_request, database_id) + set_database_id_to_request(expected_request, database_id) client._datastore_api.rollback.assert_called_once_with(request=expected_request) @@ -527,7 +527,7 @@ def _make_begin_request(project, read_only=False, read_time=None, database=None) "project_id": project, "transaction_options": expected_options, } - _set_database_id_to_request(request, database) + set_database_id_to_request(request, database) return request