api: fix push image unittests

BUG=none
TEST=./run_tests

Change-Id: I8f8f6bfc3c9bcf6773da2f502de5b102d572ec79
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/4832161
Auto-Submit: Jack Neus <jackneus@google.com>
Commit-Queue: Jack Neus <jackneus@google.com>
Tested-by: Jack Neus <jackneus@google.com>
Reviewed-by: Greg Edelston <gredelston@google.com>
Commit-Queue: Benjamin Shai <bshai@google.com>
Reviewed-by: Benjamin Shai <bshai@google.com>
diff --git a/api/controller/image_unittest.py b/api/controller/image_unittest.py
index 5d45a17..92a5cf3 100644
--- a/api/controller/image_unittest.py
+++ b/api/controller/image_unittest.py
@@ -634,7 +634,7 @@
         self.assertFalse(output_proto.success)
 
 
-class PushImageTest(api_config.ApiConfigMixin):
+class PushImageTest(cros_test_lib.MockTestCase, api_config.ApiConfigMixin):
     """Push image test."""
 
     def _GetRequest(
@@ -658,7 +658,7 @@
         )
 
     def _GetResponse(self):
-        return image_pb2.PushImageRequest()
+        return image_pb2.PushImageResponse()
 
     @mock.patch.object(pushimage, "PushImage", return_value={})
     def testValidateOnly(self, MockPushImage):
@@ -676,7 +676,7 @@
             ]
         )
         rc = image_controller.PushImage(
-            req, self.NewResponse(), self.validate_only_config
+            req, self._GetResponse(), self.validate_only_config
         )
         MockPushImage.assert_not_called()
         self.assertEqual(rc, controller.RETURN_CODE_VALID_INPUT)
@@ -775,7 +775,7 @@
             ),
         )
 
-    def testCallSucceeds(self, _):
+    def testCallSucceeds(self):
         """Check that a (dry run) call is made successfully."""
         request = self._GetRequest(sign_types=[common_pb2.IMAGE_TYPE_RECOVERY])
         rc = image_controller.PushImage(