Build API: Support binary serialization.

BUG=chromium:1032573
TEST=run_tests
TEST=manually ran endpoints
TEST=cq

Change-Id: I59c401d228f81a52c28d80e9db9718d41a255015
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2078962
Reviewed-by: LaMont Jones <lamontjones@chromium.org>
Reviewed-by: Michael Mortensen <mmortensen@google.com>
Commit-Queue: Alex Klein <saklein@chromium.org>
Tested-by: Alex Klein <saklein@chromium.org>
diff --git a/api/router_unittest.py b/api/router_unittest.py
index a045b7e..32942f9 100644
--- a/api/router_unittest.py
+++ b/api/router_unittest.py
@@ -13,6 +13,7 @@
 from google.protobuf import json_format
 
 from chromite.api import api_config
+from chromite.api import message_util
 from chromite.api import router
 from chromite.api.gen.chromite.api import build_api_test_pb2
 from chromite.lib import chroot_lib
@@ -27,38 +28,130 @@
 class RouterTest(cros_test_lib.RunCommandTempDirTestCase,
                  api_config.ApiConfigMixin):
   """Test Router functionality."""
-  _INPUT_JSON_TEMPLATE = '{"id":"Input ID", "chroot":{"path": "%s"}}'
 
   def setUp(self):
     self.router = router.Router()
     self.router.Register(build_api_test_pb2)
 
-    self.input_file = os.path.join(self.tempdir, 'input.json')
-    self.output_file = os.path.join(self.tempdir, 'output.json')
-
     self.chroot_dir = os.path.join(self.tempdir, 'chroot')
     chroot_tmp = os.path.join(self.chroot_dir, 'tmp')
     # Make the tmp dir for the re-exec inside chroot input/output files.
     osutils.SafeMakedirs(chroot_tmp)
 
-    osutils.WriteFile(self.input_file,
-                      self._INPUT_JSON_TEMPLATE % self.chroot_dir)
-    osutils.WriteFile(self.output_file, '{}')
+    # Build the input/output/config paths we'll be using in the tests.
+    self.json_input_file = os.path.join(self.tempdir, 'input.json')
+    self.json_output_file = os.path.join(self.tempdir, 'output.json')
+    self.json_config_file = os.path.join(self.tempdir, 'config.json')
+    self.binary_input_file = os.path.join(self.tempdir, 'input.bin')
+    self.binary_output_file = os.path.join(self.tempdir, 'output.bin')
+    self.binary_config_file = os.path.join(self.tempdir, 'config.bin')
+
+    # The message handlers for the respective files.
+    self.json_input_handler = message_util.get_message_handler(
+        self.json_input_file, message_util.FORMAT_JSON)
+    self.json_output_handler = message_util.get_message_handler(
+        self.json_output_file, message_util.FORMAT_JSON)
+    self.json_config_handler = message_util.get_message_handler(
+        self.json_config_file, message_util.FORMAT_JSON)
+    self.binary_input_handler = message_util.get_message_handler(
+        self.binary_input_file, message_util.FORMAT_BINARY)
+    self.binary_output_handler = message_util.get_message_handler(
+        self.binary_output_file, message_util.FORMAT_BINARY)
+    self.binary_config_handler = message_util.get_message_handler(
+        self.binary_config_file, message_util.FORMAT_BINARY)
+
+    # Build an input message to use.
+    self.expected_id = 'input id'
+    input_msg = build_api_test_pb2.TestRequestMessage()
+    input_msg.id = self.expected_id
+    input_msg.chroot.path = self.chroot_dir
+
+    # Write out base input and config messages.
+    osutils.WriteFile(self.json_input_file,
+                      json_format.MessageToJson(input_msg))
+    osutils.WriteFile(
+        self.binary_input_file, input_msg.SerializeToString(), mode='wb')
+
+    config_msg = self.api_config.get_proto()
+    osutils.WriteFile(self.json_config_file,
+                      json_format.MessageToJson(config_msg))
+    osutils.WriteFile(
+        self.binary_config_file, config_msg.SerializeToString(), mode='wb')
 
     self.subprocess_tempdir = os.path.join(self.chroot_dir, 'tempdir')
     osutils.SafeMakedirs(self.subprocess_tempdir)
 
-  def testInputOutputMethod(self):
-    """Test input/output handling."""
+  def testJsonInputOutputMethod(self):
+    """Test json input/output handling."""
     def impl(input_msg, output_msg, config):
       self.assertIsInstance(input_msg, build_api_test_pb2.TestRequestMessage)
       self.assertIsInstance(output_msg, build_api_test_pb2.TestResultMessage)
       self.assertIsInstance(config, api_config.ApiConfig)
+      self.assertEqual(config, self.api_config)
 
     self.PatchObject(self.router, '_GetMethod', return_value=impl)
 
-    self.router.Route('chromite.api.TestApiService', 'InputOutputMethod',
-                      self.input_file, self.output_file, self.api_config)
+    self.router.Route(
+        'chromite.api.TestApiService',
+        'InputOutputMethod',
+        self.api_config,
+        self.json_input_handler,
+        [self.json_output_handler],
+        self.json_config_handler)
+
+  def testBinaryInputOutputMethod(self):
+    """Test binary input/output handling."""
+    def impl(input_msg, output_msg, config):
+      self.assertIsInstance(input_msg, build_api_test_pb2.TestRequestMessage)
+      self.assertIsInstance(output_msg, build_api_test_pb2.TestResultMessage)
+      self.assertIsInstance(config, api_config.ApiConfig)
+      self.assertEqual(config, self.api_config)
+
+    self.PatchObject(self.router, '_GetMethod', return_value=impl)
+
+    self.router.Route(
+        'chromite.api.TestApiService',
+        'InputOutputMethod',
+        self.api_config,
+        self.binary_input_handler,
+        [self.binary_output_handler],
+        self.binary_config_handler)
+
+  def testMultipleOutputHandling(self):
+    """Test multiple output handling."""
+    expected_result = 'Success!'
+
+    def impl(input_msg, output_msg, config):
+      self.assertIsInstance(input_msg, build_api_test_pb2.TestRequestMessage)
+      self.assertIsInstance(output_msg, build_api_test_pb2.TestResultMessage)
+      self.assertIsInstance(config, api_config.ApiConfig)
+      self.assertEqual(config, self.api_config)
+      # Set the property on the output to test against.
+      output_msg.result = expected_result
+
+    self.PatchObject(self.router, '_GetMethod', return_value=impl)
+
+    self.router.Route(
+        'chromite.api.TestApiService',
+        'InputOutputMethod',
+        self.api_config,
+        self.binary_input_handler,
+        [self.binary_output_handler, self.json_output_handler],
+        self.binary_config_handler)
+
+    # Make sure it did write out all the expected files.
+    self.assertExists(self.binary_output_file)
+    self.assertExists(self.json_output_file)
+
+    # Parse the output files back into a message.
+    binary_msg = build_api_test_pb2.TestResultMessage()
+    json_msg = build_api_test_pb2.TestResultMessage()
+    self.binary_output_handler.read_into(binary_msg)
+    self.json_output_handler.read_into(json_msg)
+
+    # Make sure the parsed messages have the expected content.
+    self.assertEqual(binary_msg.result, expected_result)
+    self.assertEqual(json_msg.result, expected_result)
 
   def testRenameMethod(self):
     """Test implementation name config."""
@@ -69,7 +162,8 @@
     self.PatchObject(self.router, '_GetMethod', side_effect=_GetMethod)
 
     self.router.Route('chromite.api.TestApiService', 'RenamedMethod',
-                      self.input_file, self.output_file, self.api_config)
+                      self.api_config, self.binary_input_handler,
+                      [self.binary_output_handler], self.binary_config_handler)
 
   def _mock_callable(self, expect_called):
     """Helper to create the implementation mock to test chroot assertions.
@@ -87,13 +181,14 @@
 
     return impl
 
-  def _writeChrootCallOutput(self, content='{}'):
+  def _writeChrootCallOutput(self, content='{}', mode='w'):
     def impl(*_args, **_kwargs):
       """Side effect for inside-chroot calls to the API."""
       osutils.WriteFile(
           os.path.join(self.subprocess_tempdir,
                        router.Router.REEXEC_OUTPUT_FILE),
-          content)
+          content,
+          mode=mode)
 
     return impl
 
@@ -103,8 +198,9 @@
                      return_value=self._mock_callable(expect_called=True))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=True)
     self.router.Route('chromite.api.InsideChrootApiService',
-                      'InsideServiceInsideMethod', self.input_file,
-                      self.output_file, self.api_config)
+                      'InsideServiceInsideMethod', self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
   def testInsideServiceOutsideMethodOutsideChroot(self):
     """Test the outside method override works as expected."""
@@ -112,21 +208,22 @@
                      return_value=self._mock_callable(expect_called=True))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
     self.router.Route('chromite.api.InsideChrootApiService',
-                      'InsideServiceOutsideMethod', self.input_file,
-                      self.output_file, self.api_config)
+                      'InsideServiceOutsideMethod', self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
   def testInsideServiceInsideMethodOutsideChroot(self):
     """Test calling an inside method from outside the chroot."""
     self.PatchObject(self.router, '_GetMethod',
                      return_value=self._mock_callable(expect_called=False))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
-    self.rc.SetDefaultCmdResult(side_effect=self._writeChrootCallOutput())
 
     service = 'chromite.api.InsideChrootApiService'
     method = 'InsideServiceInsideMethod'
     service_method = '%s/%s' % (service, method)
-    self.router.Route(service, method, self.input_file, self.output_file,
-                      self.api_config)
+    self.router.Route(service, method, self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
     self.assertCommandContains(['build_api', service_method], enter_chroot=True)
 
@@ -137,8 +234,9 @@
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=True)
     with self.assertRaises(cros_build_lib.DieSystemExit):
       self.router.Route('chromite.api.InsideChrootApiService',
-                        'InsideServiceOutsideMethod', self.input_file,
-                        self.output_file, self.api_config)
+                        'InsideServiceOutsideMethod', self.api_config,
+                        self.binary_input_handler, [self.binary_output_handler],
+                        self.binary_config_handler)
 
   def testOutsideServiceOutsideMethodOutsideChroot(self):
     """Test outside/outside/outside works correctly."""
@@ -146,8 +244,9 @@
                      return_value=self._mock_callable(expect_called=True))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
     self.router.Route('chromite.api.OutsideChrootApiService',
-                      'OutsideServiceOutsideMethod', self.input_file,
-                      self.output_file, self.api_config)
+                      'OutsideServiceOutsideMethod', self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
   def testOutsideServiceInsideMethodInsideChroot(self):
     """Test the inside method assertion override works properly."""
@@ -155,27 +254,27 @@
                      return_value=self._mock_callable(expect_called=True))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=True)
     self.router.Route('chromite.api.OutsideChrootApiService',
-                      'OutsideServiceInsideMethod', self.input_file,
-                      self.output_file, self.api_config)
+                      'OutsideServiceInsideMethod', self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
   def testOutsideServiceInsideMethodOutsideChroot(self):
     """Test calling an inside override method from outside the chroot."""
     self.PatchObject(self.router, '_GetMethod',
                      return_value=self._mock_callable(expect_called=False))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
-    self.rc.SetDefaultCmdResult(side_effect=self._writeChrootCallOutput())
 
     service = 'chromite.api.OutsideChrootApiService'
     method = 'OutsideServiceInsideMethod'
     service_method = '%s/%s' % (service, method)
-    self.router.Route(service, method, self.input_file, self.output_file,
-                      self.api_config)
+    self.router.Route(service, method, self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
     self.assertCommandContains(['build_api', service_method], enter_chroot=True)
 
   def testReexecNonemptyOutput(self):
     """Test calling an inside chroot method that produced output."""
-    osutils.WriteFile(self.output_file, '')
     self.PatchObject(self.router, '_GetMethod',
                      return_value=self._mock_callable(expect_called=False))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
@@ -187,30 +286,62 @@
     tempdir.tempdir = self.subprocess_tempdir
     self.PatchObject(chroot_lib.Chroot, 'tempdir', return_value=tempdir)
 
+    expected_output_msg = build_api_test_pb2.TestResultMessage()
+    expected_output_msg.result = 'foo'
+
+    # Set the command side effect to write out our expected output to the
+    # output file for the inside the chroot reexecution of the endpoint.
+    # This lets us make sure the logic moving everything out works as intended.
     self.rc.SetDefaultCmdResult(
-        side_effect=self._writeChrootCallOutput(content='{"result": "foo"}'))
+        side_effect=self._writeChrootCallOutput(
+            content=expected_output_msg.SerializeToString(), mode='wb'))
 
     service = 'chromite.api.OutsideChrootApiService'
     method = 'OutsideServiceInsideMethod'
     service_method = '%s/%s' % (service, method)
-    self.router.Route(service, method, self.input_file, self.output_file,
-                      self.api_config)
+    self.router.Route(service, method, self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
     self.assertCommandContains(['build_api', service_method], enter_chroot=True)
 
     # It should be writing the result out to our output file.
-    expected = build_api_test_pb2.TestResultMessage()
-    json_format.Parse('{"result": "foo"}', expected)
-    result = build_api_test_pb2.TestResultMessage()
-    json_format.Parse(osutils.ReadFile(self.output_file), result)
-    self.assertEqual(expected, result)
+    output_msg = build_api_test_pb2.TestResultMessage()
+    self.binary_output_handler.read_into(output_msg)
+    self.assertEqual(expected_output_msg, output_msg)
 
     tempdir.tempdir = original
     del tempdir
 
   def testReexecEmptyOutput(self):
     """Test calling an inside chroot method that produced no output."""
-    osutils.WriteFile(self.output_file, '')
+    self.PatchObject(self.router, '_GetMethod',
+                     return_value=self._mock_callable(expect_called=False))
+    self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
+    expected_output_msg = build_api_test_pb2.TestResultMessage()
+
+    # Set the command side effect to write out our expected output to the
+    # output file for the inside the chroot reexecution of the endpoint.
+    # This lets us make sure the logic moving everything out works as intended.
+    self.rc.SetDefaultCmdResult(
+        side_effect=self._writeChrootCallOutput(
+            content=expected_output_msg.SerializeToString(), mode='wb'))
+
+    service = 'chromite.api.OutsideChrootApiService'
+    method = 'OutsideServiceInsideMethod'
+    service_method = '%s/%s' % (service, method)
+    self.router.Route(service, method, self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
+
+    self.assertCommandContains(['build_api', service_method], enter_chroot=True)
+
+    output_msg = build_api_test_pb2.TestResultMessage()
+    self.binary_output_handler.read_into(output_msg)
+    self.assertEqual(expected_output_msg, output_msg)
+
+  def testReexecNoOutput(self):
+    """Test calling an inside chroot method that produced no output."""
     self.PatchObject(self.router, '_GetMethod',
                      return_value=self._mock_callable(expect_called=False))
     self.PatchObject(cros_build_lib, 'IsInsideChroot', return_value=False)
@@ -219,12 +350,16 @@
     service = 'chromite.api.OutsideChrootApiService'
     method = 'OutsideServiceInsideMethod'
     service_method = '%s/%s' % (service, method)
-    self.router.Route(service, method, self.input_file, self.output_file,
-                      self.api_config)
+    self.router.Route(service, method, self.api_config,
+                      self.binary_input_handler, [self.binary_output_handler],
+                      self.binary_config_handler)
 
     self.assertCommandContains(['build_api', service_method], enter_chroot=True)
-    # It should be writing the empty message out.
-    self.assertFileContents(self.output_file, '{}')
+
+    output_msg = build_api_test_pb2.TestResultMessage()
+    empty_msg = build_api_test_pb2.TestResultMessage()
+    self.binary_output_handler.read_into(output_msg)
+    self.assertEqual(empty_msg, output_msg)
 
   def testInvalidService(self):
     """Test invalid service call."""
@@ -232,8 +367,9 @@
     method = 'OutsideServiceInsideMethod'
 
     with self.assertRaises(router.UnknownServiceError):
-      self.router.Route(service, method, self.input_file, self.output_file,
-                        self.api_config)
+      self.router.Route(service, method, self.api_config,
+                        self.binary_input_handler, [self.binary_output_handler],
+                        self.binary_config_handler)
 
   def testInvalidMethod(self):
     """Test invalid method call."""
@@ -241,5 +377,6 @@
     method = 'DoesNotExist'
 
     with self.assertRaises(router.UnknownMethodError):
-      self.router.Route(service, method, self.input_file, self.output_file,
-                        self.api_config)
+      self.router.Route(service, method, self.api_config,
+                        self.binary_input_handler, [self.binary_output_handler],
+                        self.binary_config_handler)