From 2eab49df98f11e32d3cb49de25ba43705ecf3219 Mon Sep 17 00:00:00 2001 From: yash30201 <54198301+yash30201@users.noreply.github.com> Date: Thu, 28 Mar 2024 11:23:24 +0000 Subject: [PATCH] Update beginTransaction and allocateIds rpc + unit + snippet tests --- .../DatastoreOperationRefreshTrait.php | 68 +++++++++++++++ Datastore/src/Operation.php | 42 ++++++++- .../tests/Snippet/DatastoreClientTest.php | 62 +++++++------ .../Snippet/DatastoreSessionHandlerTest.php | 33 ++++--- .../tests/Snippet/ReadOnlyTransactionTest.php | 22 ++--- Datastore/tests/Snippet/TransactionTest.php | 52 +++++------ Datastore/tests/Unit/DatastoreClientTest.php | 86 ++++++++----------- Datastore/tests/Unit/OperationTest.php | 64 +++++++------- Datastore/tests/Unit/TransactionTest.php | 24 +++--- 9 files changed, 277 insertions(+), 176 deletions(-) diff --git a/Core/src/Testing/DatastoreOperationRefreshTrait.php b/Core/src/Testing/DatastoreOperationRefreshTrait.php index bc9e943bdb2c..ef845e1b1140 100644 --- a/Core/src/Testing/DatastoreOperationRefreshTrait.php +++ b/Core/src/Testing/DatastoreOperationRefreshTrait.php @@ -22,6 +22,8 @@ use Google\Cloud\Datastore\Connection\ConnectionInterface; use Google\Cloud\Datastore\EntityMapper; use Google\Cloud\Datastore\Operation; +use Google\Cloud\Datastore\V1\Client\DatastoreClient; +use Prophecy\Argument; /** * Refresh Datastore operation class @@ -88,4 +90,70 @@ public function refreshOperation($stub, ConnectionInterface $connection, Request return $stub; } + + /** + * Helper method for Unit and Snippet test classes. This mocks the + * $requestHandler class property present in the Test Class with + * given arguments. + * + * @param string $methodName The method name to mock in RequestHandler::sendRequest + * @param array $params The parameters to look for in the + * array equivalent of rpc request. + * @param mixed $returnValue The value to be returned by sendRequest mock. + * @param null|int $shouldBeCalledTimes Adds a shouldBeCalled prophecy. Defaults to `null`, implying nothing is added. + * [ + * `0` => `shouldBeCalled`, + * Non zero positive integer $x => `shouldBeCalledTimes($x)` + * ] + */ + private function mockSendRequest($methodName, $params, $returnValue, $shouldBeCalledTimes = null) + { + if (isset($this->serializer)) { + $serializer = $this->serializer; + } else { + $serializer = new Serializer([], [ + 'google.protobuf.Value' => function ($v) { + return $this->flattenValue($v); + }, + 'google.protobuf.Timestamp' => function ($v) { + return $this->formatTimestampFromApi($v); + } + ], [], [ + 'google.protobuf.Timestamp' => function ($v) { + if (is_string($v)) { + $dt = new \DateTime($v); + return ['seconds' => $dt->format('U')]; + } + return $v; + } + ]); + } + + $prophecy = $this->requestHandler->sendRequest( + DatastoreClient::class, + $methodName, + Argument::that(function ($arg) use ($methodName, $params, $serializer) { + $requestName = ucfirst($methodName . 'Request'); + $x = explode('\\', get_class($arg)); + $argName = end($x); + + if ($requestName != $argName) { + return false; + } + $data = $serializer->encodeMessage($arg); + return array_merge($params, $data) == $data; + }), + Argument::cetera() + ); + + if (!is_null($shouldBeCalledTimes)) { + if ($shouldBeCalledTimes == 0) { + $prophecy->shouldBeCalled(); + } else { + $prophecy->shouldBeCalledTimes($shouldBeCalledTimes); + } + } + + $prophecy->willReturn($returnValue); + } } diff --git a/Datastore/src/Operation.php b/Datastore/src/Operation.php index 6daa367f005d..a3ae3a5285a9 100644 --- a/Datastore/src/Operation.php +++ b/Datastore/src/Operation.php @@ -28,7 +28,11 @@ use Google\Cloud\Datastore\Query\AggregationQueryResult; use Google\Cloud\Datastore\Query\Query; use Google\Cloud\Datastore\Query\QueryInterface; +use Google\Cloud\Datastore\V1\AllocateIdsRequest; +use Google\Cloud\Datastore\V1\BeginTransactionRequest; +use Google\Cloud\Datastore\V1\Client\DatastoreClient; use Google\Cloud\Datastore\V1\QueryResultBatch\MoreResultsType; +use Google\Cloud\Datastore\V1\TransactionOptions; /** * Run lookups and queries and commit changes. @@ -304,11 +308,30 @@ public function beginTransaction($transactionOptions, array $options = []) $transactionOptions['readOnly'] ); } - $res = $this->connection->beginTransaction($options + [ + + array_walk($transactionOptions, function (&$item) { + if ($item instanceof \stdClass) { + $item = []; + } + }); + + list($data, $optionalArgs) = $this->splitOptionalArgs($options); + + $data += [ 'projectId' => $this->projectId, 'databaseId' => $this->databaseId, 'transactionOptions' => $transactionOptions, - ]); + ]; + + $data = $this->convertDataToProtos($data, ['transactionOptions' => TransactionOptions::class]); + $request = $this->serializer->decodeMessage(new BeginTransactionRequest(), $data); + + $res = $this->requestHandler->sendRequest( + DatastoreClient::class, + 'beginTransaction', + $request, + $optionalArgs + ); return $res['transaction']; } @@ -349,11 +372,22 @@ public function allocateIds(array $keys, array $options = []) $serviceKeys[] = $key->keyObject(); } - $res = $this->connection->allocateIds($options + [ + list($data, $optionalArgs) = $this->splitOptionalArgs($options); + + $data += [ 'projectId' => $this->projectId, 'databaseId' => $this->databaseId, 'keys' => $serviceKeys, - ]); + ]; + + $request = $this->serializer->decodeMessage(new AllocateIdsRequest(), $data); + + $res = $this->requestHandler->sendRequest( + DatastoreClient::class, + 'allocateIds', + $request, + $optionalArgs + ); if (isset($res['keys'])) { foreach ($res['keys'] as $index => $key) { diff --git a/Datastore/tests/Snippet/DatastoreClientTest.php b/Datastore/tests/Snippet/DatastoreClientTest.php index ce010f2ca8a4..9e61b90fc4ad 100644 --- a/Datastore/tests/Snippet/DatastoreClientTest.php +++ b/Datastore/tests/Snippet/DatastoreClientTest.php @@ -17,6 +17,7 @@ namespace Google\Cloud\Datastore\Tests\Snippet; +use Google\Cloud\Core\ApiHelperTrait; use Google\Cloud\Core\Int64; use Google\Cloud\Core\RequestHandler; use Google\Cloud\Core\Testing\DatastoreOperationRefreshTrait; @@ -36,6 +37,7 @@ use Google\Cloud\Datastore\Query\QueryInterface; use Google\Cloud\Datastore\ReadOnlyTransaction; use Google\Cloud\Datastore\Transaction; +use Google\Cloud\Datastore\V1\Client\DatastoreClient as DatastoreGapicClient; use Prophecy\Argument; use Prophecy\PhpUnit\ProphecyTrait; @@ -44,6 +46,7 @@ */ class DatastoreClientTest extends SnippetTestCase { + use ApiHelperTrait; use DatastoreOperationRefreshTrait; use ProphecyTrait; @@ -58,6 +61,7 @@ class DatastoreClientTest extends SnippetTestCase public function setUp(): void { $this->connection = $this->prophesize(ConnectionInterface::class); + $this->client = TestHelpers::stub(DatastoreClient::class, [], ['operation']); $this->key = new Key('my-awesome-project', [ @@ -351,11 +355,11 @@ public function testTransaction() $snippet = $this->snippetFromMethod(DatastoreClient::class, 'transaction'); $snippet->addLocal('datastore', $this->client); - $this->connection->beginTransaction($this->validateTransactionOptions('readWrite')) - ->shouldBeCalled() - ->willReturn([ - 'transaction' => 'foo' - ]); + $this->mockSendRequest( + 'beginTransaction', + ['transactionOptions' => ['readWrite' => []]], + ['transaction' => 'foo'] + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -369,11 +373,11 @@ public function testReadOnlyTransaction() { $snippet = $this->snippetFromMethod(DatastoreClient::class, 'readOnlyTransaction'); $snippet->addLocal('datastore', $this->client); - $this->connection->beginTransaction($this->validateTransactionOptions('readOnly')) - ->shouldBeCalled() - ->willReturn([ - 'transaction' => 'foo' - ]); + $this->mockSendRequest( + 'beginTransaction', + ['transactionOptions' => ['readOnly' => []]], + ['transaction' => 'foo'] + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT ]); @@ -770,28 +774,30 @@ public function testRunAggregationQuery() private function allocateIdsConnectionMock() { - $this->connection->allocateIds(Argument::any()) - ->shouldBeCalled() - ->willReturn([ - 'keys' => [ - [ - 'path' => [ - [ - 'kind' => 'Person', - 'id' => '4682475895' - ] + $this->requestHandler->sendRequest( + DatastoreGapicClient::class, + 'allocateIds', + Argument::cetera() + )->shouldBeCalled()->willReturn([ + 'keys' => [ + [ + 'path' => [ + [ + 'kind' => 'Person', + 'id' => '4682475895' ] - ], - [ - 'path' => [ - [ - 'kind' => 'Person', - 'id' => '4682475896' - ] + ] + ], + [ + 'path' => [ + [ + 'kind' => 'Person', + 'id' => '4682475896' ] ] ] - ]); + ] + ]); } private function validateTransactionOptions($type, array $options = []) diff --git a/Datastore/tests/Snippet/DatastoreSessionHandlerTest.php b/Datastore/tests/Snippet/DatastoreSessionHandlerTest.php index 8e09b3fe00ae..ada15639d968 100644 --- a/Datastore/tests/Snippet/DatastoreSessionHandlerTest.php +++ b/Datastore/tests/Snippet/DatastoreSessionHandlerTest.php @@ -17,6 +17,7 @@ namespace Google\Cloud\Datastore\Tests\Snippet; +use Google\Cloud\Core\ApiHelperTrait; use Google\Cloud\Core\RequestHandler; use Google\Cloud\Core\Testing\DatastoreOperationRefreshTrait; use Google\Cloud\Core\Testing\Snippet\SnippetTestCase; @@ -36,8 +37,10 @@ class DatastoreSessionHandlerTest extends SnippetTestCase { use DatastoreOperationRefreshTrait; use ProphecyTrait; + use ApiHelperTrait; const TRANSACTION = 'transaction-id'; + const PROJECT = 'example-project'; private $connection; private $client; @@ -78,12 +81,11 @@ public function testClass() }) ))->shouldBeCalled()->willReturn([]); - $this->connection->beginTransaction(Argument::withEntry( - 'transactionOptions', - ['readWrite' => (object) []] - ))->shouldBeCalled()->willReturn([ - 'transaction' => self::TRANSACTION, - ]); + $this->mockSendRequest( + 'beginTransaction', + ['transactionOptions' => ['readWrite' => []]], + ['transaction' => self::TRANSACTION] + ); $this->connection->commit(Argument::allOf( Argument::withEntry('transaction', self::TRANSACTION), @@ -99,7 +101,9 @@ public function testClass() }) ))->shouldBeCalled()->willReturn([]); - $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal()); + $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ + 'projectId' => self::PROJECT + ]); $snippet->addLocal('datastore', $this->client); $res = $snippet->invoke(); @@ -124,12 +128,11 @@ public function testClassErrorHandler() }) ))->shouldBeCalled()->willReturn([]); - $this->connection->beginTransaction(Argument::withEntry( - 'transactionOptions', - ['readWrite' => (object) []] - ))->shouldBeCalled()->willReturn([ - 'transaction' => self::TRANSACTION, - ]); + $this->mockSendRequest( + 'beginTransaction', + ['transactionOptions' => ['readWrite' => []]], + ['transaction' => self::TRANSACTION] + ); $this->connection->commit(Argument::any()) ->shouldBeCalled() @@ -137,7 +140,9 @@ public function testClassErrorHandler() trigger_error('oops!', E_USER_WARNING); }); - $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal()); + $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ + 'projectId' => self::PROJECT + ]); $snippet->addLocal('datastore', $this->client); $res = $snippet->invoke(); diff --git a/Datastore/tests/Snippet/ReadOnlyTransactionTest.php b/Datastore/tests/Snippet/ReadOnlyTransactionTest.php index 022a49225b6d..b4bd041eae12 100644 --- a/Datastore/tests/Snippet/ReadOnlyTransactionTest.php +++ b/Datastore/tests/Snippet/ReadOnlyTransactionTest.php @@ -104,11 +104,12 @@ public function setUp(): void public function testClass() { - $this->connection->beginTransaction(Argument::any()) - ->shouldBeCalled() - ->willReturn([ - 'transaction' => 'foo' - ]); + $this->mockSendRequest( + 'beginTransaction', + [], + ['transaction' => 'foo'], + 0 + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -124,11 +125,12 @@ public function testClass() public function testClassRollback() { - $this->connection->beginTransaction(Argument::any()) - ->shouldBeCalled() - ->willReturn([ - 'transaction' => 'foo' - ]); + $this->mockSendRequest( + 'beginTransaction', + [], + ['transaction' => 'foo'], + 0 + ); $this->connection->lookup(Argument::any()) ->shouldBeCalled() ->willReturn([]); diff --git a/Datastore/tests/Snippet/TransactionTest.php b/Datastore/tests/Snippet/TransactionTest.php index bc14c638611d..3f00669260a5 100644 --- a/Datastore/tests/Snippet/TransactionTest.php +++ b/Datastore/tests/Snippet/TransactionTest.php @@ -31,6 +31,7 @@ use Google\Cloud\Datastore\Query\AggregationQuery; use Google\Cloud\Datastore\Query\QueryInterface; use Google\Cloud\Datastore\Transaction; +use Google\Cloud\Datastore\V1\Client\DatastoreClient as DatastoreGapicClient; use Prophecy\Argument; use Prophecy\PhpUnit\ProphecyTrait; @@ -105,11 +106,12 @@ public function setUp(): void public function testClass() { - $this->connection->beginTransaction(Argument::any()) - ->shouldBeCalled() - ->willReturn([ - 'transaction' => 'foo' - ]); + $this->mockSendRequest( + 'beginTransaction', + [], + ['transaction' => 'foo'], + 0 + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -506,29 +508,31 @@ public function testRollback() private function allocateIdsConnectionMock() { - $this->connection->allocateIds(Argument::any()) - ->shouldBeCalled() - ->willReturn([ - 'keys' => [ - [ - 'path' => [ - [ - 'kind' => 'Person', - 'id' => '4682475895' - ] + $this->requestHandler->sendRequest( + DatastoreGapicClient::class, + 'allocateIds', + Argument::cetera() + )->shouldBeCalled()->willReturn([ + 'keys' => [ + [ + 'path' => [ + [ + 'kind' => 'Person', + 'id' => '4682475895' ] - ], - [ - 'path' => [ - [ - 'kind' => 'Person', - 'id' => '4682475896' - ] + ] + ], + [ + 'path' => [ + [ + 'kind' => 'Person', + 'id' => '4682475896' ] ] ] - ]); + ] + ]); - return $this->connection->reveal(); + return $this->requestHandler->reveal(); } } diff --git a/Datastore/tests/Unit/DatastoreClientTest.php b/Datastore/tests/Unit/DatastoreClientTest.php index 3a2898eb4fe1..eb3d55bc3786 100644 --- a/Datastore/tests/Unit/DatastoreClientTest.php +++ b/Datastore/tests/Unit/DatastoreClientTest.php @@ -215,13 +215,12 @@ public function testAllocateId($method, $batch = false) $keyWithId = clone $key; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::withEntry('keys', [ - $key->keyObject() - ]))->shouldBeCalled()->willReturn([ - 'keys' => [ - $keyWithId->keyObject() - ] - ]); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$key->keyObject()]], + ['keys' => [$keyWithId->keyObject()]], + 0 + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -249,23 +248,14 @@ public function allocateIdProvider() */ public function testTransaction($method, $type, $key) { - $this->connection->beginTransaction(Argument::allOf( - Argument::withEntry('projectId', self::PROJECT), - // can't do direct comparisons between (object)[]. - Argument::that(function ($arg) use ($key) { - if (!($arg['transactionOptions'][$key] instanceof \stdClass)) { - return false; - } - - if ((array) $arg['transactionOptions'][$key]) { - return false; - } - - return true; - }) - ))->shouldBeCalled()->willReturn([ - 'transaction' => self::TRANSACTION - ]); + $this->mockSendRequest( + 'beginTransaction', + [ + 'projectId' => self::PROJECT, + 'transactionOptions' => ['readTime' => []] + ], + ['transaction' => self::TRANSACTION] + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -280,16 +270,14 @@ public function testTransaction($method, $type, $key) */ public function testTransactionWithOptions($method, $type, $key) { - $options = ['foo' => 'bar']; - - $this->connection->beginTransaction(Argument::allOf( - Argument::withEntry('projectId', self::PROJECT), - Argument::withEntry('transactionOptions', [ - $key => $options - ]) - ))->shouldBeCalled()->willReturn([ - 'transaction' => self::TRANSACTION - ]); + $this->mockSendRequest( + 'beginTransaction', + [ + 'projectId' => self::PROJECT, + 'transactionOptions' => [$key => []] + ], + ['transaction' => self::TRANSACTION] + ); // Make sure the correct transaction ID was injected. $this->connection->runQuery(Argument::withEntry('transaction', self::TRANSACTION)) @@ -300,7 +288,7 @@ public function testTransactionWithOptions($method, $type, $key) 'projectId' => self::PROJECT ]); - $res = $this->client->$method(['transactionOptions' => $options]); + $res = $this->client->$method(['transactionOptions' => []]); $this->assertInstanceOf($type, $res); iterator_to_array($res->runQuery($this->client->gqlQuery('SELECT 1=1'))); @@ -376,13 +364,11 @@ public function testMutationsWithPartialKey($method, $mutation, $key, $id) $keyWithId = clone $key; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::allOf( - Argument::withEntry('keys', [$key->keyObject()]) - ))->shouldBeCalled()->willReturn([ - 'keys' => [ - $keyWithId->keyObject() - ] - ]); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$key->keyObject()]], + ['keys' => [$keyWithId->keyObject()]] + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -405,13 +391,11 @@ public function testBatchMutationsWithPartialKey($method, $mutation, $key, $id) $keyWithId = clone $key; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::allOf( - Argument::withEntry('keys', [$key->keyObject()]) - ))->shouldBeCalled()->willReturn([ - 'keys' => [ - $keyWithId->keyObject() - ] - ]); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$key->keyObject()]], + ['keys' => [$keyWithId->keyObject()]] + ); $this->refreshOperation($this->client, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -684,7 +668,7 @@ public function testRunAggregationQuery() 'aggregateProperties' => ['property_1' => 1] ] ], - 'readTime' => (new \DateTime)->format('Y-m-d\TH:i:s') .'.000001Z' + 'readTime' => (new \DateTime())->format('Y-m-d\TH:i:s') .'.000001Z' ] ]); @@ -720,7 +704,7 @@ public function testAggregationQueryWithDifferentReturnTypes($response, $expecte 'aggregateProperties' => ['property_1' => $response] ] ], - 'readTime' => (new \DateTime)->format('Y-m-d\TH:i:s') .'.000001Z' + 'readTime' => (new \DateTime())->format('Y-m-d\TH:i:s') .'.000001Z' ] ]); diff --git a/Datastore/tests/Unit/OperationTest.php b/Datastore/tests/Unit/OperationTest.php index 7765e1238077..9782dcc59af5 100644 --- a/Datastore/tests/Unit/OperationTest.php +++ b/Datastore/tests/Unit/OperationTest.php @@ -20,6 +20,7 @@ use Google\ApiCore\Serializer; use Google\Cloud\Core\ApiHelperTrait; use Google\Cloud\Core\RequestHandler; +use Google\Cloud\Core\Testing\DatastoreOperationRefreshTrait; use Google\Cloud\Core\Testing\TestHelpers; use Google\Cloud\Datastore\Connection\ConnectionInterface; use Google\Cloud\Datastore\Entity; @@ -43,10 +44,11 @@ class OperationTest extends TestCase { use ProphecyTrait; use ApiHelperTrait; + use DatastoreOperationRefreshTrait; - const PROJECT = 'example-project'; - const NAMESPACEID = 'namespace-id'; - const DATABASEID = 'database-id'; + public const PROJECT = 'example-project'; + public const NAMESPACEID = 'namespace-id'; + public const DATABASEID = 'database-id'; private $operation; private $connection; @@ -81,7 +83,7 @@ public function setUp(): void null, new EntityMapper('foo', true, false), self::DATABASEID, - ], ['connection', 'namespaceId']); + ], ['connection', 'requestHandler', 'namespaceId']); } public function testKey() @@ -251,15 +253,14 @@ public function testAllocateIds() $id = 12345; $keyWithId = clone $key; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::withEntry('keys', [$key->keyObject()])) - ->shouldBeCalled() - ->willReturn([ - 'keys' => [ - $keyWithId->keyObject(), - ], - ]); - $this->operation->___setProperty('connection', $this->connection->reveal()); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$key->keyObject()]], + ['keys' => [$keyWithId->keyObject()]] + ); + + $this->operation->___setProperty('requestHandler', $this->requestHandler->reveal()); $res = $this->operation->allocateIds([$key]); @@ -724,15 +725,14 @@ public function testAllocateIdsToEntities() $id = 12345; $keyWithId = clone $partialKey; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::withEntry('keys', [$partialKey->keyObject()])) - ->shouldBeCalled() - ->willReturn([ - 'keys' => [ - $keyWithId->keyObject(), - ], - ]); - $this->operation->___setProperty('connection', $this->connection->reveal()); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$partialKey->keyObject()]], + ['keys' => [$keyWithId->keyObject()]] + ); + + $this->operation->___setProperty('requestHandler', $this->requestHandler->reveal()); $entities = [ $this->operation->entity($completeKey), @@ -988,11 +988,11 @@ public function testInvalidBatchType() public function testBeginTransactionWithDatabaseIdOverride() { - $this->connection - ->beginTransaction( - Argument::withEntry('databaseId', 'otherDatabaseId') - ) - ->willReturn(['transaction' => 'valid_test_transaction']); + $this->mockSendRequest( + 'beginTransaction', + ['databaseId' => 'otherDatabaseId'], + ['transaction' => 'valid_test_transaction'] + ); $transactionId = $this->operation->beginTransaction( [], @@ -1004,12 +1004,14 @@ public function testBeginTransactionWithDatabaseIdOverride() public function testAllocateIdsWithDatabaseIdOverride() { - $this->connection - ->allocateIds( - Argument::withEntry('databaseId', 'otherDatabaseId') - ) - ->shouldBeCalledTimes(1) - ->willReturn([]); + $this->mockSendRequest( + 'allocateIds', + ['databaseId' => 'otherDatabaseId'], + [], + 1 + ); + + $this->operation->___setProperty('requestHandler', $this->requestHandler->reveal()); $this->operation->allocateIds( [], diff --git a/Datastore/tests/Unit/TransactionTest.php b/Datastore/tests/Unit/TransactionTest.php index 68f4ee37f3ed..fc3a55412df1 100644 --- a/Datastore/tests/Unit/TransactionTest.php +++ b/Datastore/tests/Unit/TransactionTest.php @@ -413,13 +413,11 @@ public function testMutationsWithPartialKey($method, $mutation, $key, $id) $keyWithId = clone $key; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::allOf( - Argument::withEntry('keys', [$key->keyObject()]) - ))->shouldBeCalled()->willReturn([ - 'keys' => [ - $keyWithId->keyObject() - ] - ]); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$key->keyObject()]], + ['keys' => [$keyWithId->keyObject()]] + ); $this->refreshOperation($this->transaction, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT @@ -445,13 +443,11 @@ public function testBatchMutationsWithPartialKey($method, $mutation, $key, $id) $keyWithId = clone $key; $keyWithId->setLastElementIdentifier($id); - $this->connection->allocateIds(Argument::allOf( - Argument::withEntry('keys', [$key->keyObject()]) - ))->shouldBeCalled()->willReturn([ - 'keys' => [ - $keyWithId->keyObject() - ] - ]); + $this->mockSendRequest( + 'allocateIds', + ['keys' => [$key->keyObject()]], + ['keys' => [$keyWithId->keyObject()]] + ); $this->refreshOperation($this->transaction, $this->connection->reveal(), $this->requestHandler->reveal(), [ 'projectId' => self::PROJECT