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<string, mixed> $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