diff --git a/composer.json b/composer.json
index 20469fa..e5ba59d 100644
--- a/composer.json
+++ b/composer.json
@@ -5,7 +5,7 @@
     "minimum-stability": "stable",
     "require": {
         "google/cloud-core": "^1.14",
-        "google/gax": "^0.35"
+        "google/gax": "^0.36"
     },
     "require-dev": {
         "phpunit/phpunit": "^4.8|^5.0",
diff --git a/tests/Unit/V1/ImageAnnotatorClientExtensionTest.php b/tests/Unit/V1/ImageAnnotatorClientExtensionTest.php
index 9b1f0d1..ee44b5a 100644
--- a/tests/Unit/V1/ImageAnnotatorClientExtensionTest.php
+++ b/tests/Unit/V1/ImageAnnotatorClientExtensionTest.php
@@ -115,7 +115,12 @@ public function testAnnotateImageWithImageContext()
         $this->transport->startUnaryCall( Argument::allOf(
                     Argument::type(Call::class),
                     Argument::which('getMethod', 'google.cloud.vision.v1.ImageAnnotator/BatchAnnotateImages'),
-                    Argument::which('getMessage', $expectedMessage)
+                    Argument::that(function($arg) use ($expectedMessage) {
+                        if ($arg->getMessage()->serializeToString() === $expectedMessage->serializeToString()) {
+                            return true;
+                        }
+                        return false;
+                    })
                 ),
                 Argument::type('array')
             )
@@ -160,7 +165,12 @@ public function testDetectionMethod($methodName, $featureType, $image)
                 Argument::allOf(
                     Argument::type(Call::class),
                     Argument::which('getMethod', 'google.cloud.vision.v1.ImageAnnotator/BatchAnnotateImages'),
-                    Argument::which('getMessage', $expectedMessage)
+                    Argument::that(function($arg) use ($expectedMessage) {
+                        if ($arg->getMessage()->serializeToString() === $expectedMessage->serializeToString()) {
+                            return true;
+                        }
+                        return false;
+                    })
                 ),
                 Argument::type('array')
             )