# Copyright 2018 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Tests for the build_api script covering the base Build API functionality."""

import os
from typing import Callable

from chromite.third_party.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
from chromite.lib import constants
from chromite.lib import cros_build_lib
from chromite.lib import cros_test_lib
from chromite.lib import osutils


# The executable we expect to be called when re-execing build_api
# inside the SDK.
_BUILD_API_INSIDE = (
    constants.CHROOT_SOURCE_ROOT / constants.CHROMITE_BIN_SUBDIR / "build_api"
)


class RouterTest(
    cros_test_lib.RunCommandTempDirTestCase, api_config.ApiConfigMixin
):
    """Test Router functionality."""

    def setUp(self) -> None:
        self.router = router.Router()
        self.router.Register(build_api_test_pb2)

        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=False)

        self.chroot_dir = self.tempdir / "chroot"
        self.out_dir = self.tempdir / "out"
        chroot = chroot_lib.Chroot(
            path=self.chroot_dir,
            out_path=self.out_dir,
        )
        # Make sure basic path structures exist.
        osutils.SafeMakedirs(self.chroot_dir)
        osutils.SafeMakedirs(self.out_dir)
        # Make the tmp dir for the re-exec inside chroot input/output files.
        osutils.SafeMakedirs(chroot.tmp)

        # 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 = str(self.chroot_dir)
        input_msg.chroot.out_path = str(self.out_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 testJsonInputOutputMethod(self) -> None:
        """Test json input/output handling."""

        def impl(input_msg, output_msg, config) -> None:
            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.json_input_handler,
            [self.json_output_handler],
            self.json_config_handler,
        )

    def testBinaryInputOutputMethod(self) -> None:
        """Test binary input/output handling."""

        def impl(input_msg, output_msg, config) -> None:
            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) -> None:
        """Test multiple output handling."""
        expected_result = "Success!"

        def impl(input_msg, output_msg, config) -> None:
            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) -> None:
        """Test implementation name config."""

        def _GetMethod(_, method_name):
            self.assertEqual("CorrectName", method_name)
            return lambda x, y, z: None

        self.PatchObject(self.router, "_GetMethod", side_effect=_GetMethod)

        self.router.Route(
            "chromite.api.TestApiService",
            "RenamedMethod",
            self.api_config,
            self.binary_input_handler,
            [self.binary_output_handler],
            self.binary_config_handler,
        )

    def _mock_callable(self, expect_called: bool) -> Callable:
        """Helper to create the implementation mock to test chroot assertions.

        Args:
            expect_called: Whether the implementation should be called.
                When False, an assertion will fail if it is called.

        Returns:
            The implementation.
        """

        def impl(_input_msg, _output_msg, _config) -> None:
            self.assertTrue(
                expect_called, "The implementation should not have been called."
            )

        return impl

    def _writeChrootCallOutput(self, content="{}", mode="w"):
        def impl(*_args, **_kwargs) -> None:
            """Side effect for inside-chroot calls to the API."""
            osutils.WriteFile(
                os.path.join(
                    self.subprocess_tempdir, router.Router.REEXEC_OUTPUT_FILE
                ),
                content,
                mode=mode,
            )

        return impl

    def testInsideServiceInsideMethodInsideChroot(self) -> None:
        """Test inside/inside/inside works correctly."""
        self.PatchObject(
            self.router,
            "_GetMethod",
            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.api_config,
            self.binary_input_handler,
            [self.binary_output_handler],
            self.binary_config_handler,
        )

    def testInsideServiceOutsideMethodOutsideChroot(self) -> None:
        """Test the outside method override works as expected."""
        self.PatchObject(
            self.router,
            "_GetMethod",
            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.api_config,
            self.binary_input_handler,
            [self.binary_output_handler],
            self.binary_config_handler,
        )

    def testInsideServiceInsideMethodOutsideChroot(self) -> None:
        """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)

        service = "chromite.api.InsideChrootApiService"
        method = "InsideServiceInsideMethod"
        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_INSIDE, service_method], enter_chroot=True
        )

    def testInsideServiceOutsideMethodInsideChroot(self) -> None:
        """Test inside chroot for outside method raises an error."""
        self.PatchObject(
            self.router,
            "_GetMethod",
            return_value=self._mock_callable(expect_called=False),
        )
        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=True)
        with self.assertRaises(cros_build_lib.DieSystemExit):
            self.router.Route(
                "chromite.api.InsideChrootApiService",
                "InsideServiceOutsideMethod",
                self.api_config,
                self.binary_input_handler,
                [self.binary_output_handler],
                self.binary_config_handler,
            )

    def testOutsideServiceOutsideMethodOutsideChroot(self) -> None:
        """Test outside/outside/outside works correctly."""
        self.PatchObject(
            self.router,
            "_GetMethod",
            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.api_config,
            self.binary_input_handler,
            [self.binary_output_handler],
            self.binary_config_handler,
        )

    def testOutsideServiceInsideMethodInsideChroot(self) -> None:
        """Test the inside method assertion override works properly."""
        self.PatchObject(
            self.router,
            "_GetMethod",
            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.api_config,
            self.binary_input_handler,
            [self.binary_output_handler],
            self.binary_config_handler,
        )

    def testOutsideServiceInsideMethodOutsideChroot(self) -> None:
        """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)

        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_INSIDE, service_method], enter_chroot=True
        )

    def testReexecNonemptyOutput(self) -> None:
        """Test calling an inside chroot method that produced output."""
        self.PatchObject(
            self.router,
            "_GetMethod",
            return_value=self._mock_callable(expect_called=False),
        )
        self.PatchObject(cros_build_lib, "IsInsideChroot", return_value=False)

        # Patch the chroot tempdir method to return a tempdir with our
        # subprocess tempdir so the output file will be in the expected
        # location.
        tempdir = osutils.TempDir()
        original = tempdir.tempdir
        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=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_INSIDE, service_method], enter_chroot=True
        )

        # It should be writing the result out to our output file.
        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) -> None:
        """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)
        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_INSIDE, 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) -> None:
        """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)
        self.rc.SetDefaultCmdResult(returncode=1)

        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_INSIDE, service_method], enter_chroot=True
        )

        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 test_tot_service_tot_method(self) -> None:
        """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) -> None:
        """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) -> None:
        """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_branched_call_no_config(self) -> None:
        """Re-execute branched BAPI with no config passed."""
        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"

        config_handler = message_util.get_message_handler(
            None, message_util.FORMAT_BINARY
        )
        self.router.Route(
            service,
            method,
            self.api_config,
            self.binary_input_handler,
            [self.binary_output_handler],
            config_handler,
        )

        self.assertCommandContains(
            [self.tempdir / "bin" / "build_api", f"{service}/{method}"]
        )
        # Branched BAPI calls don't parse and rewrite the messages, so when no
        # config is given, an empty one isn't written out like with SDK reexecs.
        self.assertCommandContains(
            ["--config-binary", "--config-json"], expected=False
        )

    def test_tot_service_branched_method_inside(self) -> None:
        """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) -> None:
        """Test invalid service call."""
        service = "chromite.api.DoesNotExist"
        method = "OutsideServiceInsideMethod"

        with self.assertRaises(router.UnknownServiceError):
            self.router.Route(
                service,
                method,
                self.api_config,
                self.binary_input_handler,
                [self.binary_output_handler],
                self.binary_config_handler,
            )

    def testInvalidMethod(self) -> None:
        """Test invalid method call."""
        service = "chromite.api.OutsideChrootApiService"
        method = "DoesNotExist"

        with self.assertRaises(router.UnknownMethodError):
            self.router.Route(
                service,
                method,
                self.api_config,
                self.binary_input_handler,
                [self.binary_output_handler],
                self.binary_config_handler,
            )

    def testListVisibility(self) -> None:
        """Test visibility options."""
        service = "HiddenService"
        method = "HiddenMethod"

        for current in self.router.ListMethods():
            self.assertNotIn(service, current)
            self.assertNotIn(method, current)