#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import unittest
from collections import defaultdict
from typing import List, Optional, Tuple
from unittest.mock import MagicMock, call, patch

from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus
from fbpcp.service.mpc import MPCInstanceStatus, MPCParty, MPCService
from fbpcp.service.onedocker import OneDockerService
from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
from fbpcs.data_processing.lift_id_combiner.lift_id_spine_combiner_cpp import (
    CppLiftIdSpineCombinerService,
)
from fbpcs.data_processing.sharding.sharding_cpp import CppShardingService
from fbpcs.onedocker_binary_config import OneDockerBinaryConfig
from fbpcs.onedocker_binary_names import OneDockerBinaryNames
from fbpcs.onedocker_service_config import OneDockerServiceConfig
from fbpcs.pcf.tests.async_utils import to_sync
from fbpcs.pid.entity.pid_instance import (
    PIDInstance,
    PIDInstanceStatus,
    PIDProtocol,
    PIDRole,
)
from fbpcs.pid.service.pid_service.pid import PIDService
from fbpcs.private_computation.entity.private_computation_instance import (
    PrivateComputationGameType,
    PrivateComputationInstance,
    PrivateComputationInstanceStatus,
    PrivateComputationRole,
    UnionedPCInstance,
)
from fbpcs.private_computation.entity.private_computation_stage_type import (
    PrivateComputationStageType,
)
from fbpcs.private_computation.repository.private_computation_game import GameNames
from fbpcs.private_computation.service.errors import (
    PrivateComputationServiceValidationError,
)
from fbpcs.private_computation.service.private_computation import (
    PrivateComputationService,
    NUM_NEW_SHARDS_PER_FILE,
    DEFAULT_K_ANONYMITY_THRESHOLD,
)
from fbpcs.private_computation.service.private_computation_stage_service import (
    PrivateComputationStageService,
)

# TODO T94666166: libfb won't work in OSS
from libfb.py.asyncio.mock import AsyncMock
from libfb.py.testutil import data_provider

from fbpcs.private_computation.service.utils import (
    create_and_start_mpc_instance,
    gen_mpc_game_args_to_retry,
    map_private_computation_role_to_mpc_party,
    DEFAULT_CONTAINER_TIMEOUT_IN_SEC,
)


def _get_valid_stages_data() -> List[Tuple[PrivateComputationStageType]]:
    return [
        (PrivateComputationStageType.ID_MATCH,),
        (PrivateComputationStageType.COMPUTE,),
        (PrivateComputationStageType.AGGREGATE,),
        (PrivateComputationStageType.POST_PROCESSING_HANDLERS,),
    ]


class TestPrivateComputationService(unittest.TestCase):
    def setUp(self):
        container_svc_patcher = patch("fbpcp.service.container_aws.AWSContainerService")
        storage_svc_patcher = patch("fbpcp.service.storage_s3.S3StorageService")
        mpc_instance_repo_patcher = patch(
            "fbpcs.common.repository.mpc_instance_local.LocalMPCInstanceRepository"
        )
        pid_instance_repo_patcher = patch(
            "fbpcs.pid.repository.pid_instance_local.LocalPIDInstanceRepository"
        )
        private_computation_instance_repo_patcher = patch(
            "fbpcs.private_computation.repository.private_computation_instance_local.LocalPrivateComputationInstanceRepository"
        )
        mpc_game_svc_patcher = patch("fbpcp.service.mpc_game.MPCGameService")
        container_svc = container_svc_patcher.start()
        storage_svc = storage_svc_patcher.start()
        mpc_instance_repository = mpc_instance_repo_patcher.start()
        pid_instance_repository = pid_instance_repo_patcher.start()
        private_computation_instance_repository = (
            private_computation_instance_repo_patcher.start()
        )
        mpc_game_svc = mpc_game_svc_patcher.start()

        for patcher in (
            container_svc_patcher,
            storage_svc_patcher,
            mpc_instance_repo_patcher,
            pid_instance_repo_patcher,
            private_computation_instance_repo_patcher,
            mpc_game_svc_patcher,
        ):
            self.addCleanup(patcher.stop)

        self.onedocker_service_config = OneDockerServiceConfig(
            task_definition="test_task_definition",
        )

        self.onedocker_binary_config_map = defaultdict(
            lambda: OneDockerBinaryConfig(
                tmp_directory="/test_tmp_directory/", binary_version="latest"
            )
        )

        self.onedocker_service = OneDockerService(
            container_svc, self.onedocker_service_config.task_definition
        )

        self.mpc_service = MPCService(
            container_svc=container_svc,
            instance_repository=mpc_instance_repository,
            task_definition="test_task_definition",
            mpc_game_svc=mpc_game_svc,
        )

        self.pid_service = PIDService(
            instance_repository=pid_instance_repository,
            storage_svc=storage_svc,
            onedocker_svc=self.onedocker_service,
            onedocker_binary_config_map=self.onedocker_binary_config_map,
        )

        self.private_computation_service = PrivateComputationService(
            instance_repository=private_computation_instance_repository,
            mpc_svc=self.mpc_service,
            pid_svc=self.pid_service,
            onedocker_svc=self.onedocker_service,
            onedocker_binary_config_map=self.onedocker_binary_config_map,
        )

        self.test_private_computation_id = "test_private_computation_id"
        self.test_num_containers = 2
        self.test_input_path = "in_path"
        self.test_output_dir = "out_dir"
        self.test_game_type = PrivateComputationGameType.LIFT
        self.test_concurrency = 1

    def test_create_instance(self):
        test_role = PrivateComputationRole.PUBLISHER
        self.private_computation_service.create_instance(
            instance_id=self.test_private_computation_id,
            role=test_role,
            game_type=self.test_game_type,
            input_path=self.test_input_path,
            output_dir=self.test_output_dir,
            num_pid_containers=self.test_num_containers,
            num_mpc_containers=self.test_num_containers,
            concurrency=self.test_concurrency,
            num_files_per_mpc_container=NUM_NEW_SHARDS_PER_FILE,
        )
        # check instance_repository.create is called with the correct arguments
        self.private_computation_service.instance_repository.create.assert_called()
        args = self.private_computation_service.instance_repository.create.call_args[0][
            0
        ]
        self.assertEqual(self.test_private_computation_id, args.instance_id)
        self.assertEqual(test_role, args.role)
        self.assertEqual(PrivateComputationInstanceStatus.CREATED, args.status)

    def test_update_instance(self):
        test_pid_id = self.test_private_computation_id + "_id_match"
        test_pid_protocol = PIDProtocol.UNION_PID
        test_pid_role = PIDRole.PUBLISHER
        test_input_path = "pid_in"
        test_output_path = "pid_out"
        # create one PID instance to be put into PrivateComputationInstance
        pid_instance = PIDInstance(
            instance_id=test_pid_id,
            protocol=test_pid_protocol,
            pid_role=test_pid_role,
            num_shards=self.test_num_containers,
            input_path=test_input_path,
            output_path=test_output_path,
            status=PIDInstanceStatus.STARTED,
        )

        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.ID_MATCHING_STARTED,
            instances=[pid_instance],
        )

        updated_pid_instance = pid_instance
        updated_pid_instance.status = PIDInstanceStatus.COMPLETED
        self.private_computation_service.pid_svc.update_instance = MagicMock(
            return_value=updated_pid_instance
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )

        # call update on the PrivateComputationInstance
        updated_instance = self.private_computation_service.update_instance(
            instance_id=self.test_private_computation_id
        )

        # check update instance called on the right pid instance
        self.private_computation_service.pid_svc.update_instance.assert_called()
        self.assertEqual(
            test_pid_id,
            self.private_computation_service.pid_svc.update_instance.call_args[0][0],
        )

        # check update instance called on the right private lift instance
        self.private_computation_service.instance_repository.update.assert_called()
        self.assertEqual(
            private_computation_instance,
            self.private_computation_service.instance_repository.update.call_args[0][0],
        )

        # check updated_instance has new status
        self.assertEqual(
            PrivateComputationInstanceStatus.ID_MATCHING_COMPLETED,
            updated_instance.status,
        )

        # create one MPC instance to be put into PrivateComputationInstance
        test_mpc_id = "test_mpc_id"
        mpc_instance = PCSMPCInstance.create_instance(
            instance_id=test_mpc_id,
            game_name=GameNames.LIFT.value,
            mpc_party=MPCParty.SERVER,
            num_workers=2,
        )

        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.COMPUTATION_STARTED,
            instances=[mpc_instance],
        )

        updated_mpc_instance = mpc_instance
        updated_mpc_instance.status = MPCInstanceStatus.COMPLETED
        self.private_computation_service.mpc_svc.update_instance = MagicMock(
            return_value=updated_mpc_instance
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )
        # call update on the PrivateComputationInstance
        updated_instance = self.private_computation_service.update_instance(
            instance_id=self.test_private_computation_id
        )

        # check update instance called on the right mpc instance
        self.private_computation_service.mpc_svc.update_instance.assert_called()
        self.assertEqual(
            test_mpc_id,
            self.private_computation_service.mpc_svc.update_instance.call_args[0][0],
        )

        # check update instance called on the right private lift instance
        self.private_computation_service.instance_repository.update.assert_called()
        self.assertEqual(
            private_computation_instance,
            self.private_computation_service.instance_repository.update.call_args[0][0],
        )

        # check updated_instance has new status
        self.assertEqual(
            PrivateComputationInstanceStatus.COMPUTATION_COMPLETED,
            updated_instance.status,
        )

    @staticmethod
    def _get_dummy_stage_svc(
        stage_type: PrivateComputationStageType,
    ) -> PrivateComputationStageService:
        """create a DummyTestStageService class and instantiate an instance of it"""

        return type(
            "DummyTestStageService",
            (PrivateComputationStageService,),
            {
                "run_async": AsyncMock(
                    # run_async will return whatever pc_instance privatelift.run_stage passes it
                    side_effect=lambda pc_instance, *args, **kwargs: pc_instance
                ),
                "stage_type": stage_type,
            },
        )()

    @data_provider(_get_valid_stages_data)
    def test_run_stage_correct_stage_order(
        self, stage_type: PrivateComputationStageType
    ) -> None:
        """
        tests that run_stage runs stage_svc when the stage_svc is the next stage in the sequence
        """
        ################# PREVIOUS STAGE COMPLETED OR RETRY #######################
        stage_svc = self._get_dummy_stage_svc(stage_type)
        for status in (
            stage_type.previous_stage.completed_status,
            stage_type.failed_status,
        ):
            pl_instance = self.create_sample_instance(status=status)
            self.private_computation_service.instance_repository.read = MagicMock(
                return_value=pl_instance
            )

            pl_instance = self.private_computation_service.run_stage(
                pl_instance.instance_id, stage_svc
            )
            self.assertEqual(pl_instance.status, stage_type.start_status)

    @data_provider(_get_valid_stages_data)
    def test_run_stage_status_already_started(
        self, stage_type: PrivateComputationStageType
    ) -> None:
        """
        tests that run_stage does not run stage_svc when the instance status is already started
        """
        ################# CURRENT STAGE STATUS NOT VALID #######################
        stage_svc = self._get_dummy_stage_svc(stage_type)
        pl_instance = self.create_sample_instance(status=stage_type.start_status)

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=pl_instance
        )

        with self.assertRaises(ValueError):
            pl_instance = self.private_computation_service.run_stage(
                pl_instance.instance_id, stage_svc
            )

    @data_provider(_get_valid_stages_data)
    def test_run_stage_out_of_order_with_dry_run(
        self, stage_type: PrivateComputationStageType
    ) -> None:
        """
        tests that run_stage runs stage_svc out of order when dry run is passed
        """
        ################ STAGE OUT OF ORDER WITH DRY RUN #####################
        stage_svc = self._get_dummy_stage_svc(stage_type)
        pl_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.UNKNOWN
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=pl_instance
        )

        pl_instance = self.private_computation_service.run_stage(
            pl_instance.instance_id, stage_svc, dry_run=True
        )
        self.assertEqual(pl_instance.status, stage_type.start_status)

    @data_provider(_get_valid_stages_data)
    def test_run_stage_out_of_order_without_dry_run(
        self, stage_type: PrivateComputationStageType
    ) -> None:
        """
        tests that run_stage does not run stage_svc out of order when dry run is not passed
        """
        ####################### STAGE OUT OF ORDER NO DRY RUN ############################
        stage_svc = self._get_dummy_stage_svc(stage_type)
        pl_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.UNKNOWN
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=pl_instance
        )

        with self.assertRaises(ValueError):
            pl_instance = self.private_computation_service.run_stage(
                pl_instance.instance_id, stage_svc, dry_run=False
            )

    @data_provider(_get_valid_stages_data)
    def test_run_stage_partner_no_server_ips(
        self, stage_type: PrivateComputationStageType
    ) -> None:
        """
        tests that run_stage does not if role is partner and no server ips are specified
        """
        ####################### PARTNER NO SERVER IPS ############################
        stage_svc = self._get_dummy_stage_svc(stage_type)
        pl_instance = self.create_sample_instance(
            status=stage_type.previous_stage.completed_status,
            role=PrivateComputationRole.PARTNER,
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=pl_instance
        )

        with self.assertRaises(ValueError):
            pl_instance = self.private_computation_service.run_stage(
                pl_instance.instance_id, stage_svc
            )

    @data_provider(_get_valid_stages_data)
    def test_run_stage_fails(self, stage_type: PrivateComputationStageType) -> None:
        """
        tests that statuses are set properly when a run fails
        """
        ######################### STAGE FAILS ####################################
        stage_svc = self._get_dummy_stage_svc(stage_type)
        pl_instance = self.create_sample_instance(
            status=stage_type.previous_stage.completed_status
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=pl_instance
        )

        # create a custom exception class to make sure we have a unique exception for the test
        stage_failure_exception = type("TestStageFailureException", (Exception,), {})
        stage_svc.run_async = AsyncMock(side_effect=stage_failure_exception())

        with self.assertRaises(stage_failure_exception):
            pl_instance = self.private_computation_service.run_stage(
                pl_instance.instance_id, stage_svc
            )

        self.assertEqual(pl_instance.status, stage_type.failed_status)


    def test_partner_missing_server_ips(self):
        test_private_computation_id = "test_private_computation_id"

        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.ID_MATCHING_COMPLETED,
        )
        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )

        # exception because role is partner but server ips are not given
        with self.assertRaises(ValueError):
            self.private_computation_service.aggregate_shards(
                instance_id=test_private_computation_id,
            )

    @patch("fbpcp.service.mpc.MPCService")
    @patch(
        "fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance"
    )
    def test_aggregate_shards(self, mock_create_and_start_mpc_instance, mock_mpc_svc):
        # construct a private_computation_instance with an mpc_instance handling metrics computation
        test_mpc_id = self.test_private_computation_id + "_compute_metrics"
        mpc_instance = PCSMPCInstance.create_instance(
            instance_id=test_mpc_id,
            game_name=GameNames.LIFT.value,
            mpc_party=MPCParty.SERVER,
            num_workers=self.test_num_containers,
            status=MPCInstanceStatus.COMPLETED,
        )
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.COMPUTATION_COMPLETED,
            instances=[mpc_instance],
        )
        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )
        mock_mpc_svc.update_instance = MagicMock(return_value=mpc_instance)

        # call aggregate_shards
        self.private_computation_service.aggregate_shards(
            instance_id=self.test_private_computation_id,
            server_ips=["192.0.2.0", "192.0.2.1"],
        )

        test_game_args = [
            {
                "input_base_path": private_computation_instance.compute_stage_output_base_path,
                "metrics_format_type": "lift",
                "num_shards": self.test_num_containers * NUM_NEW_SHARDS_PER_FILE,
                "output_path": private_computation_instance.shard_aggregate_stage_output_path,
                "threshold": private_computation_instance.k_anonymity_threshold,
                "run_name": "",
            }
        ]
        # check a new MPC instance handling metrics aggregation was to be created
        self.assertEqual(
            GameNames.SHARD_AGGREGATOR.value,
            mock_create_and_start_mpc_instance.call_args[1]["game_name"],
        )
        self.assertEqual(
            test_game_args,
            mock_create_and_start_mpc_instance.call_args[1]["game_args"],
        )
        self.private_computation_service.instance_repository.update.assert_called()
        self.assertEqual(
            PrivateComputationInstanceStatus.AGGREGATION_STARTED,
            private_computation_instance.status,
        )

    @patch("fbpcp.service.mpc.MPCService")
    @patch(
        "fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance"
    )
    def test_aggregate_shards_rerun(
        self, mock_create_and_start_mpc_instance, mock_mpc_svc
    ):
        # construct a private_computation_instance
        test_private_computation_id = "test_private_computation_id"
        mpc_instance = PCSMPCInstance.create_instance(
            instance_id=test_private_computation_id + "_aggregate_shards",
            game_name=GameNames.SHARD_AGGREGATOR.value,
            mpc_party=MPCParty.SERVER,
            num_workers=2,
            status=MPCInstanceStatus.FAILED,
        )
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.AGGREGATION_FAILED,
            instances=[mpc_instance],
        )

        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )
        mock_mpc_svc.update_instance = MagicMock(return_value=mpc_instance)

        # call aggregate_shards
        self.private_computation_service.aggregate_shards(
            instance_id=test_private_computation_id,
            server_ips=["192.0.2.0", "192.0.2.1"],
        )

        # check that the retry counter has been incremented
        self.assertEqual(private_computation_instance.retry_counter, 1)

        # check a new MPC instance handling metrics aggregation was to be created
        self.assertEqual(2, len(private_computation_instance.instances))
        self.assertEqual(
            test_private_computation_id + "_aggregate_shards1",
            mock_create_and_start_mpc_instance.call_args[1]["instance_id"],
        )
        self.assertEqual(
            PrivateComputationInstanceStatus.AGGREGATION_STARTED,
            private_computation_instance.status,
        )

    @patch("fbpcp.service.mpc.MPCService")
    @patch(
        "fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance"
    )
    def test_aggregate_shards_dry_run(
        self, mock_create_and_start_mpc_instance, mock_mpc_svc
    ):
        # construct a private_computation_instance
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.COMPUTATION_FAILED,
        )
        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )

        # call aggregate_shards with ad-hoc input_path and num_shards
        test_format_type = "lift"
        test_game_args = [
            {
                "input_base_path": private_computation_instance.compute_stage_output_base_path,
                "num_shards": self.test_num_containers * NUM_NEW_SHARDS_PER_FILE,
                "metrics_format_type": test_format_type,
                "output_path": private_computation_instance.shard_aggregate_stage_output_path,
                "threshold": private_computation_instance.k_anonymity_threshold,
                "run_name": "",
            }
        ]
        self.private_computation_service.aggregate_shards(
            instance_id=self.test_private_computation_id,
            server_ips=["192.0.2.0", "192.0.2.1"],
            dry_run=True,
        )

        # check a new MPC instance handling metrics aggregation was to be created
        # with the overwritten input_path and num_shards
        self.assertEqual(
            GameNames.SHARD_AGGREGATOR.value,
            mock_create_and_start_mpc_instance.call_args[1]["game_name"],
        )
        self.assertEqual(
            test_game_args,
            mock_create_and_start_mpc_instance.call_args[1]["game_args"],
        )
        self.private_computation_service.instance_repository.update.assert_called()
        self.assertEqual(
            PrivateComputationInstanceStatus.AGGREGATION_STARTED,
            private_computation_instance.status,
        )

    @to_sync
    @patch("fbpcp.service.mpc.MPCService")
    async def test_create_and_start_mpc_instance(self, mock_mpc_svc):
        mock_mpc_svc.create_instance = MagicMock()
        mock_mpc_svc.start_instance_async = AsyncMock()

        instance_id = "test_instance_id"
        game_name = GameNames.LIFT.value
        mpc_party = MPCParty.CLIENT
        num_containers = 4
        input_file = "input_file"
        output_file = "output_file"
        input_directory = "input_directory"
        output_directory = "output_directory"
        server_ips = ["192.0.2.0", "192.0.2.1"]
        game_args = {
            "input_filenames": input_file,
            "input_directory": input_directory,
            "output_filenames": output_file,
            "output_directory": output_directory,
            "concurrency": 1,
        }
        binary_version = self.onedocker_binary_config_map[
            OneDockerBinaryNames.LIFT_COMPUTE.value
        ].binary_version

        await create_and_start_mpc_instance(
            mpc_svc=mock_mpc_svc,
            instance_id=instance_id,
            game_name=game_name,
            mpc_party=mpc_party,
            num_containers=num_containers,
            binary_version=binary_version,
            container_timeout=DEFAULT_CONTAINER_TIMEOUT_IN_SEC,
            server_ips=server_ips,
            game_args=game_args,
        )

        # check create_instance and start_instance were called with the right parameters
        self.assertEqual(
            call(
                instance_id=instance_id,
                game_name=game_name,
                mpc_party=mpc_party,
                num_workers=num_containers,
                game_args=game_args,
            ),
            mock_mpc_svc.create_instance.call_args,
        )

        self.assertEqual(
            call(
                instance_id=instance_id,
                server_ips=server_ips,
                timeout=DEFAULT_CONTAINER_TIMEOUT_IN_SEC,
                version=binary_version,
            ),
            mock_mpc_svc.start_instance_async.call_args,
        )

    def test_map_private_computation_role_to_mpc_party(self):
        self.assertEqual(
            MPCParty.SERVER,
            map_private_computation_role_to_mpc_party(PrivateComputationRole.PUBLISHER),
        )
        self.assertEqual(
            MPCParty.CLIENT,
            map_private_computation_role_to_mpc_party(PrivateComputationRole.PARTNER),
        )

    def test_get_status_from_stage(self):
        # Test get status from an MPC stage
        mpc_instance = PCSMPCInstance.create_instance(
            instance_id="test_mpc_id",
            game_name=GameNames.SHARD_AGGREGATOR.value,
            mpc_party=MPCParty.SERVER,
            num_workers=2,
            status=MPCInstanceStatus.FAILED,
        )
        self.assertEqual(
            PrivateComputationInstanceStatus.AGGREGATION_FAILED,
            self.private_computation_service._get_status_from_stage(mpc_instance),
        )

        # Test get status from the PID stage
        pid_instance = PIDInstance(
            instance_id="test_pid_id",
            protocol=PIDProtocol.UNION_PID,
            pid_role=PIDRole.PUBLISHER,
            num_shards=4,
            input_path="input",
            output_path="output",
            stages_containers={},
            stages_status={},
            status=PIDInstanceStatus.COMPLETED,
        )
        self.assertEqual(
            PrivateComputationInstanceStatus.ID_MATCHING_COMPLETED,
            self.private_computation_service._get_status_from_stage(pid_instance),
        )

    def test_prepare_data(self):
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.CREATED,
        )
        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )

        with patch.object(
            CppLiftIdSpineCombinerService,
            "combine_on_container_async",
        ) as mock_combine, patch.object(
            CppShardingService,
            "shard_on_container_async",
        ) as mock_shard:
            # call prepare_data
            self.private_computation_service.prepare_data(
                instance_id=self.test_private_computation_id,
                dry_run=True,
            )
            binary_config = self.onedocker_binary_config_map[
                OneDockerBinaryNames.LIFT_ID_SPINE_COMBINER.value
            ]
            mock_combine.assert_called_once_with(
                spine_path=private_computation_instance.pid_stage_output_spine_path,
                data_path=private_computation_instance.pid_stage_output_data_path,
                output_path=private_computation_instance.data_processing_output_path
                + "_combine",
                num_shards=self.test_num_containers,
                onedocker_svc=self.onedocker_service,
                binary_version=binary_config.binary_version,
                tmp_directory=binary_config.tmp_directory,
            )
            mock_shard.assert_called()

    def test_prepare_data_tasks_skipped(self):
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.COMPUTATION_FAILED,
        )
        private_computation_instance.partial_container_retry_enabled = True
        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )

        with patch.object(
            CppLiftIdSpineCombinerService,
            "combine_on_container_async",
        ) as mock_combine, patch.object(
            CppShardingService,
            "shard_on_container_async",
        ) as mock_shard:
            # call prepare_data
            self.private_computation_service.prepare_data(
                instance_id=self.test_private_computation_id,
            )
            # expect combining and sharding skipped because this private_computation_instance has
            #   status PrivateComputationInstanceStatus.COMPUTATION_FAILED, so this run
            #   is to recover from a previous compute metrics failure, meaning data
            #   preparation should have been done
            mock_combine.assert_not_called()
            mock_shard.assert_not_called()

    def test_validate_metrics_results_doesnt_match(self):
        self.private_computation_service.pid_svc.storage_svc.read = MagicMock()
        self.private_computation_service.pid_svc.storage_svc.read.side_effect = [
            '{"subGroupMetrics":[],"metrics":{"controlClicks":1,"testSpend":0,"controlImpressions":0,"testImpressions":0,"controlMatchCount":0,"testMatchCount":0,"controlNumConvSquared":0,"testNumConvSquared":0,"testValueSquared":0,"controlValue":0,"testValue":0,"testConverters":0,"testConversions":0,"testPopulation":0,"controlClickers":0,"testClickers":0,"controlReach":0,"testReach":0,"controlSpend":0,"testClicks":0,"controlValueSquared":0,"controlConverters":0,"controlConversions":0,"controlPopulation":0}}',
            '{"subGroupMetrics":[],"metrics":{"testSpend":0,"controlClicks":0,"controlImpressions":0,"testImpressions":0,"controlMatchCount":0,"testMatchCount":0,"controlNumConvSquared":0,"testNumConvSquared":0,"testValueSquared":0,"controlValue":0,"testValue":0,"testConverters":0,"testConversions":0,"testPopulation":0,"controlClickers":0,"testClickers":0,"controlReach":0,"testReach":0,"controlSpend":0,"testClicks":0,"controlValueSquared":0,"controlConverters":0,"controlConversions":0,"controlPopulation":0}}',
        ]
        with self.assertRaises(PrivateComputationServiceValidationError):
            self.private_computation_service.validate_metrics(
                instance_id="test_id",
                aggregated_result_path="aggregated_result_path",
                expected_result_path="expected_result_path",
            )

    def test_cancel_current_stage(self):
        test_mpc_id = self.test_private_computation_id + "_compute_metrics"
        test_game_name = GameNames.LIFT.value
        test_mpc_party = MPCParty.CLIENT

        # prepare the pl instance that will be read in to memory from the repository
        # at the beginning of the cancel_current_stage function
        mpc_instance_started = PCSMPCInstance.create_instance(
            instance_id=test_mpc_id,
            game_name=test_game_name,
            mpc_party=test_mpc_party,
            num_workers=self.test_num_containers,
            status=MPCInstanceStatus.STARTED,
        )
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.COMPUTATION_STARTED,
            role=PrivateComputationRole.PARTNER,
            instances=[mpc_instance_started],
        )
        self.private_computation_service.instance_repository.read = MagicMock(
            return_value=private_computation_instance
        )

        # prepare the mpc instance that's returned from mpc_service.stop_instance()
        mpc_instance_canceled = PCSMPCInstance.create_instance(
            instance_id=test_mpc_id,
            game_name=test_game_name,
            mpc_party=test_mpc_party,
            num_workers=self.test_num_containers,
            status=MPCInstanceStatus.CANCELED,
        )
        self.private_computation_service.mpc_svc.stop_instance = MagicMock(
            return_value=mpc_instance_canceled
        )
        self.private_computation_service.mpc_svc.instance_repository.read = MagicMock(
            return_value=mpc_instance_canceled
        )

        # call cancel, expect no exception
        private_computation_instance = (
            self.private_computation_service.cancel_current_stage(
                instance_id=self.test_private_computation_id,
            )
        )

        # assert the pl instance returned has the correct status
        self.assertEqual(
            PrivateComputationInstanceStatus.COMPUTATION_FAILED,
            private_computation_instance.status,
        )

    def test_gen_game_args_to_retry(self):
        test_input = "test_input_retry"
        mpc_instance = PCSMPCInstance.create_instance(
            instance_id="mpc_instance",
            game_name=GameNames.LIFT.value,
            mpc_party=MPCParty.SERVER,
            num_workers=2,
            status=MPCInstanceStatus.FAILED,
            containers=[
                ContainerInstance(
                    instance_id="container_instance_0",
                    status=ContainerInstanceStatus.FAILED,
                ),
                ContainerInstance(
                    instance_id="container_instance_1",
                    status=ContainerInstanceStatus.COMPLETED,
                ),
            ],
            game_args=[
                {
                    "input_filenames": test_input,
                },
                {
                    "input_filenames": "input_filenames",
                },
            ],
        )
        private_computation_instance = self.create_sample_instance(
            status=PrivateComputationInstanceStatus.COMPUTATION_FAILED,
            instances=[mpc_instance],
        )

        game_args = gen_mpc_game_args_to_retry(
            private_computation_instance
        )

        self.assertEqual(1, len(game_args))  # only 1 failed container
        self.assertEqual(test_input, game_args[0]["input_filenames"])

    def create_sample_instance(
        self,
        status: PrivateComputationInstanceStatus,
        role: PrivateComputationRole = PrivateComputationRole.PUBLISHER,
        instances: Optional[List[UnionedPCInstance]] = None,
    ) -> PrivateComputationInstance:
        return PrivateComputationInstance(
            instance_id=self.test_private_computation_id,
            role=role,
            instances=instances or [],
            status=status,
            status_update_ts=1600000000,
            num_pid_containers=self.test_num_containers,
            num_mpc_containers=self.test_num_containers,
            concurrency=self.test_concurrency,
            num_files_per_mpc_container=NUM_NEW_SHARDS_PER_FILE,
            game_type=PrivateComputationGameType.LIFT,
            input_path=self.test_input_path,
            output_dir=self.test_output_dir,
            fail_fast=True,
            k_anonymity_threshold=DEFAULT_K_ANONYMITY_THRESHOLD,
        )

[(296, 'libfb.py.testutil.data_provider', 'data_provider', False, 'from libfb.py.testutil import data_provider\n'), (319, 'libfb.py.testutil.data_provider', 'data_provider', False, 'from libfb.py.testutil import data_provider\n'), (339, 'libfb.py.testutil.data_provider', 'data_provider', False, 'from libfb.py.testutil import data_provider\n'), (361, 'libfb.py.testutil.data_provider', 'data_provider', False, 'from libfb.py.testutil import data_provider\n'), (383, 'libfb.py.testutil.data_provider', 'data_provider', False, 'from libfb.py.testutil import data_provider\n'), (406, 'libfb.py.testutil.data_provider', 'data_provider', False, 'from libfb.py.testutil import data_provider\n'), (449, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (450, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (503, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (504, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (549, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (550, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (599, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (77, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (78, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (79, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (82, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (85, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (88, 'unittest.mock.patch', 'patch', False, 'from unittest.mock import MagicMock, call, patch\n'), (108, 'fbpcs.onedocker_service_config.OneDockerServiceConfig', 'OneDockerServiceConfig', False, 'from fbpcs.onedocker_service_config import OneDockerServiceConfig\n'), (118, 'fbpcp.service.onedocker.OneDockerService', 'OneDockerService', False, 'from fbpcp.service.onedocker import OneDockerService\n'), (122, 'fbpcp.service.mpc.MPCService', 'MPCService', False, 'from fbpcp.service.mpc import MPCInstanceStatus, MPCParty, MPCService\n'), (129, 'fbpcs.pid.service.pid_service.pid.PIDService', 'PIDService', False, 'from fbpcs.pid.service.pid_service.pid import PIDService\n'), (136, 'fbpcs.private_computation.service.private_computation.PrivateComputationService', 'PrivateComputationService', False, 'from fbpcs.private_computation.service.private_computation import PrivateComputationService, NUM_NEW_SHARDS_PER_FILE, DEFAULT_K_ANONYMITY_THRESHOLD\n'), (180, 'fbpcs.pid.entity.pid_instance.PIDInstance', 'PIDInstance', False, 'from fbpcs.pid.entity.pid_instance import PIDInstance, PIDInstanceStatus, PIDProtocol, PIDRole\n'), (197, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (201, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (232, 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), (246, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (250, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (330, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (352, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (374, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (397, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (417, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (439, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (456, 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), (467, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (470, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (512, 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), (524, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (527, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (560, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (601, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (602, 'libfb.py.asyncio.mock.AsyncMock', 'AsyncMock', False, 'from libfb.py.asyncio.mock import AsyncMock\n'), (670, 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), (683, 'fbpcs.pid.entity.pid_instance.PIDInstance', 'PIDInstance', False, 'from fbpcs.pid.entity.pid_instance import PIDInstance, PIDInstanceStatus, PIDProtocol, PIDRole\n'), (703, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (739, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (762, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (781, 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), (793, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (798, 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), (805, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (808, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (857, 'fbpcs.private_computation.service.utils.gen_mpc_game_args_to_retry', 'gen_mpc_game_args_to_retry', False, 'from fbpcs.private_computation.service.utils import create_and_start_mpc_instance, gen_mpc_game_args_to_retry, map_private_computation_role_to_mpc_party, DEFAULT_CONTAINER_TIMEOUT_IN_SEC\n'), (870, 'fbpcs.private_computation.entity.private_computation_instance.PrivateComputationInstance', 'PrivateComputationInstance', False, 'from fbpcs.private_computation.entity.private_computation_instance import PrivateComputationGameType, PrivateComputationInstance, PrivateComputationInstanceStatus, PrivateComputationRole, UnionedPCInstance\n'), (310, 'unittest.mock.MagicMock', 'MagicMock', False, 'from unittest.mock import MagicMock, call, patch\n'), (624, 'fbpcs.private_computation.service.utils.create_and_start_mpc_instance', 'create_and_start_mpc_instance', False, 'from fbpcs.private_computation.service.utils import create_and_start_mpc_instance, gen_mpc_game_args_to_retry, map_private_computation_role_to_mpc_party, DEFAULT_CONTAINER_TIMEOUT_IN_SEC\n'), (638, 'unittest.mock.call', 'call', False, 'from unittest.mock import MagicMock, call, patch\n'), (649, 'unittest.mock.call', 'call', False, 'from unittest.mock import MagicMock, call, patch\n'), (661, 'fbpcs.private_computation.service.utils.map_private_computation_role_to_mpc_party', 'map_private_computation_role_to_mpc_party', False, 'from fbpcs.private_computation.service.utils import create_and_start_mpc_instance, gen_mpc_game_args_to_retry, map_private_computation_role_to_mpc_party, DEFAULT_CONTAINER_TIMEOUT_IN_SEC\n'), (665, 'fbpcs.private_computation.service.utils.map_private_computation_role_to_mpc_party', 'map_private_computation_role_to_mpc_party', False, 'from fbpcs.private_computation.service.utils import create_and_start_mpc_instance, gen_mpc_game_args_to_retry, map_private_computation_role_to_mpc_party, DEFAULT_CONTAINER_TIMEOUT_IN_SEC\n'), (707, 'unittest.mock.patch.object', 'patch.object', False, 'from unittest.mock import MagicMock, call, patch\n'), (710, 'unittest.mock.patch.object', 'patch.object', False, 'from unittest.mock import MagicMock, call, patch\n'), (743, 'unittest.mock.patch.object', 'patch.object', False, 'from unittest.mock import MagicMock, call, patch\n'), (746, 'unittest.mock.patch.object', 'patch.object', False, 'from unittest.mock import MagicMock, call, patch\n'), (113, 'fbpcs.onedocker_binary_config.OneDockerBinaryConfig', 'OneDockerBinaryConfig', False, 'from fbpcs.onedocker_binary_config import OneDockerBinaryConfig\n'), (288, 'libfb.py.asyncio.mock.AsyncMock', 'AsyncMock', False, 'from libfb.py.asyncio.mock import AsyncMock\n'), (834, 'fbpcp.entity.container_instance.ContainerInstance', 'ContainerInstance', False, 'from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus\n'), (838, 'fbpcp.entity.container_instance.ContainerInstance', 'ContainerInstance', False, 'from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus\n')]