BAPI: Re-exec on branched BAPI logic.

Initial implementation re-executing BAPI calls on the branched
BAPI when configured to do so, but without the ability to examine
or modify the requests and responses, or the branched BAPI itself.

BUG=None
TEST=manual, run_tests

Change-Id: Ie9e75db9a7ccac662cca63d869bf8167fffb5039
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/4482527
Reviewed-by: Benjamin Shai <bshai@google.com>
Tested-by: Alex Klein <saklein@chromium.org>
Commit-Queue: Alex Klein <saklein@chromium.org>
Commit-Queue: Lizzy Presland <zland@google.com>
Reviewed-by: Lizzy Presland <zland@google.com>
Auto-Submit: Alex Klein <saklein@chromium.org>
diff --git a/api/router_unittest.py b/api/router_unittest.py
index 316c0fe..4877941 100644
--- a/api/router_unittest.py
+++ b/api/router_unittest.py
@@ -14,6 +14,7 @@
 from chromite.api import router
 from chromite.api.gen.chromite.api import build_api_test_pb2
 from chromite.lib import chroot_lib
+from chromite.lib import constants
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_test_lib
 from chromite.lib import osutils
@@ -485,6 +486,101 @@
         self.binary_output_handler.read_into(output_msg)
         self.assertEqual(empty_msg, output_msg)
 
+    def test_tot_service_tot_method(self):
+        """Test no re-exec for ToT->ToT."""
+        self.PatchObject(
+            self.router,
+            "_GetMethod",
+            return_value=self._mock_callable(expect_called=True),
+        )
+        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=False)
+        self.PatchObject(constants, "IS_BRANCHED_CHROMITE", new=False)
+
+        self.router.Route(
+            "chromite.api.TotExecutionService",
+            "TotServiceTotMethod",
+            self.api_config,
+            self.binary_input_handler,
+            [self.binary_output_handler],
+            self.binary_config_handler,
+        )
+
+    def test_tot_service_tot_method_inside(self):
+        """Test error raised when method runs ToT & inside the SDK."""
+        self.PatchObject(
+            self.router,
+            "_GetMethod",
+            return_value=self._mock_callable(expect_called=False),
+        )
+        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=False)
+        self.PatchObject(constants, "IS_BRANCHED_CHROMITE", new=False)
+
+        with self.assertRaises(router.TotSdkError):
+            self.router.Route(
+                "chromite.api.TotExecutionService",
+                "TotServiceTotMethodInside",
+                self.api_config,
+                self.binary_input_handler,
+                [self.binary_output_handler],
+                self.binary_config_handler,
+            )
+
+    def test_tot_service_branched_method(self):
+        """Re-execute branched BAPI."""
+        self.PatchObject(
+            self.router,
+            "_GetMethod",
+            return_value=self._mock_callable(expect_called=False),
+        )
+        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=False)
+        self.PatchObject(constants, "IS_BRANCHED_CHROMITE", new=False)
+        self.PatchObject(constants, "BRANCHED_CHROMITE_DIR", new=self.tempdir)
+
+        service = "chromite.api.TotExecutionService"
+        method = "TotServiceBranchedMethod"
+        self.router.Route(
+            service,
+            method,
+            self.api_config,
+            self.binary_input_handler,
+            [self.binary_output_handler],
+            self.binary_config_handler,
+        )
+
+        self.assertCommandContains(
+            [self.tempdir / "bin" / "build_api", f"{service}/{method}"]
+        )
+
+    def test_tot_service_branched_method_inside(self):
+        """Re-execute branched BAPI.
+
+        Chroot handling delegated to the branched BAPI, so should be identical
+        to the branched method case above.
+        """
+        self.PatchObject(
+            self.router,
+            "_GetMethod",
+            return_value=self._mock_callable(expect_called=False),
+        )
+        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=False)
+        self.PatchObject(constants, "IS_BRANCHED_CHROMITE", new=False)
+        self.PatchObject(constants, "BRANCHED_CHROMITE_DIR", new=self.tempdir)
+
+        service = "chromite.api.TotExecutionService"
+        method = "TotServiceBranchedMethodInside"
+        self.router.Route(
+            service,
+            method,
+            self.api_config,
+            self.binary_input_handler,
+            [self.binary_output_handler],
+            self.binary_config_handler,
+        )
+
+        self.assertCommandContains(
+            [self.tempdir / "bin" / "build_api", f"{service}/{method}"]
+        )
+
     def testInvalidService(self):
         """Test invalid service call."""
         service = "chromite.api.DoesNotExist"