#!/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,
        )

[((11751, 11788), 'libfb.py.testutil.data_provider', 'data_provider', (['_get_valid_stages_data'], {}), '(_get_valid_stages_data)\n', (11764, 11788), False, 'from libfb.py.testutil import data_provider\n'), ((12717, 12754), 'libfb.py.testutil.data_provider', 'data_provider', (['_get_valid_stages_data'], {}), '(_get_valid_stages_data)\n', (12730, 12754), False, 'from libfb.py.testutil import data_provider\n'), ((13525, 13562), 'libfb.py.testutil.data_provider', 'data_provider', (['_get_valid_stages_data'], {}), '(_get_valid_stages_data)\n', (13538, 13562), False, 'from libfb.py.testutil import data_provider\n'), ((14385, 14422), 'libfb.py.testutil.data_provider', 'data_provider', (['_get_valid_stages_data'], {}), '(_get_valid_stages_data)\n', (14398, 14422), False, 'from libfb.py.testutil import data_provider\n'), ((15259, 15296), 'libfb.py.testutil.data_provider', 'data_provider', (['_get_valid_stages_data'], {}), '(_get_valid_stages_data)\n', (15272, 15296), False, 'from libfb.py.testutil import data_provider\n'), ((16152, 16189), 'libfb.py.testutil.data_provider', 'data_provider', (['_get_valid_stages_data'], {}), '(_get_valid_stages_data)\n', (16165, 16189), False, 'from libfb.py.testutil import data_provider\n'), ((17950, 17987), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.mpc.MPCService"""'], {}), "('fbpcp.service.mpc.MPCService')\n", (17955, 17987), False, 'from unittest.mock import MagicMock, call, patch\n'), ((17993, 18095), 'unittest.mock.patch', 'patch', (['"""fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance"""'], {}), "(\n    'fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance'\n    )\n", (17998, 18095), False, 'from unittest.mock import MagicMock, call, patch\n'), ((20384, 20421), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.mpc.MPCService"""'], {}), "('fbpcp.service.mpc.MPCService')\n", (20389, 20421), False, 'from unittest.mock import MagicMock, call, patch\n'), ((20427, 20529), 'unittest.mock.patch', 'patch', (['"""fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance"""'], {}), "(\n    'fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance'\n    )\n", (20432, 20529), False, 'from unittest.mock import MagicMock, call, patch\n'), ((22303, 22340), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.mpc.MPCService"""'], {}), "('fbpcp.service.mpc.MPCService')\n", (22308, 22340), False, 'from unittest.mock import MagicMock, call, patch\n'), ((22346, 22448), 'unittest.mock.patch', 'patch', (['"""fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance"""'], {}), "(\n    'fbpcs.private_computation.service.private_computation.create_and_start_mpc_instance'\n    )\n", (22351, 22448), False, 'from unittest.mock import MagicMock, call, patch\n'), ((24411, 24448), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.mpc.MPCService"""'], {}), "('fbpcp.service.mpc.MPCService')\n", (24416, 24448), False, 'from unittest.mock import MagicMock, call, patch\n'), ((2797, 2853), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.container_aws.AWSContainerService"""'], {}), "('fbpcp.service.container_aws.AWSContainerService')\n", (2802, 2853), False, 'from unittest.mock import MagicMock, call, patch\n'), ((2884, 2934), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.storage_s3.S3StorageService"""'], {}), "('fbpcp.service.storage_s3.S3StorageService')\n", (2889, 2934), False, 'from unittest.mock import MagicMock, call, patch\n'), ((2971, 3049), 'unittest.mock.patch', 'patch', (['"""fbpcs.common.repository.mpc_instance_local.LocalMPCInstanceRepository"""'], {}), "('fbpcs.common.repository.mpc_instance_local.LocalMPCInstanceRepository')\n", (2976, 3049), False, 'from unittest.mock import MagicMock, call, patch\n'), ((3108, 3183), 'unittest.mock.patch', 'patch', (['"""fbpcs.pid.repository.pid_instance_local.LocalPIDInstanceRepository"""'], {}), "('fbpcs.pid.repository.pid_instance_local.LocalPIDInstanceRepository')\n", (3113, 3183), False, 'from unittest.mock import MagicMock, call, patch\n'), ((3258, 3390), 'unittest.mock.patch', 'patch', (['"""fbpcs.private_computation.repository.private_computation_instance_local.LocalPrivateComputationInstanceRepository"""'], {}), "(\n    'fbpcs.private_computation.repository.private_computation_instance_local.LocalPrivateComputationInstanceRepository'\n    )\n", (3263, 3390), False, 'from unittest.mock import MagicMock, call, patch\n'), ((3434, 3480), 'unittest.mock.patch', 'patch', (['"""fbpcp.service.mpc_game.MPCGameService"""'], {}), "('fbpcp.service.mpc_game.MPCGameService')\n", (3439, 3480), False, 'from unittest.mock import MagicMock, call, patch\n'), ((4252, 4314), 'fbpcs.onedocker_service_config.OneDockerServiceConfig', 'OneDockerServiceConfig', ([], {'task_definition': '"""test_task_definition"""'}), "(task_definition='test_task_definition')\n", (4274, 4314), False, 'from fbpcs.onedocker_service_config import OneDockerServiceConfig\n'), ((4574, 4652), 'fbpcp.service.onedocker.OneDockerService', 'OneDockerService', (['container_svc', 'self.onedocker_service_config.task_definition'], {}), '(container_svc, self.onedocker_service_config.task_definition)\n', (4590, 4652), False, 'from fbpcp.service.onedocker import OneDockerService\n'), ((4703, 4863), 'fbpcp.service.mpc.MPCService', 'MPCService', ([], {'container_svc': 'container_svc', 'instance_repository': 'mpc_instance_repository', 'task_definition': '"""test_task_definition"""', 'mpc_game_svc': 'mpc_game_svc'}), "(container_svc=container_svc, instance_repository=\n    mpc_instance_repository, task_definition='test_task_definition',\n    mpc_game_svc=mpc_game_svc)\n", (4713, 4863), False, 'from fbpcp.service.mpc import MPCInstanceStatus, MPCParty, MPCService\n'), ((4942, 5131), 'fbpcs.pid.service.pid_service.pid.PIDService', 'PIDService', ([], {'instance_repository': 'pid_instance_repository', 'storage_svc': 'storage_svc', 'onedocker_svc': 'self.onedocker_service', 'onedocker_binary_config_map': 'self.onedocker_binary_config_map'}), '(instance_repository=pid_instance_repository, storage_svc=\n    storage_svc, onedocker_svc=self.onedocker_service,\n    onedocker_binary_config_map=self.onedocker_binary_config_map)\n', (4952, 5131), False, 'from fbpcs.pid.service.pid_service.pid import PIDService\n'), ((5226, 5477), 'fbpcs.private_computation.service.private_computation.PrivateComputationService', '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'}), '(instance_repository=\n    private_computation_instance_repository, mpc_svc=self.mpc_service,\n    pid_svc=self.pid_service, onedocker_svc=self.onedocker_service,\n    onedocker_binary_config_map=self.onedocker_binary_config_map)\n', (5251, 5477), False, 'from fbpcs.private_computation.service.private_computation import PrivateComputationService, NUM_NEW_SHARDS_PER_FILE, DEFAULT_K_ANONYMITY_THRESHOLD\n'), ((7278, 7510), 'fbpcs.pid.entity.pid_instance.PIDInstance', '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'}), '(instance_id=test_pid_id, protocol=test_pid_protocol, pid_role=\n    test_pid_role, num_shards=self.test_num_containers, input_path=\n    test_input_path, output_path=test_output_path, status=PIDInstanceStatus\n    .STARTED)\n', (7289, 7510), False, 'from fbpcs.pid.entity.pid_instance import PIDInstance, PIDInstanceStatus, PIDProtocol, PIDRole\n'), ((7959, 8003), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'updated_pid_instance'}), '(return_value=updated_pid_instance)\n', (7968, 8003), False, 'from unittest.mock import MagicMock, call, patch\n'), ((8095, 8147), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (8104, 8147), False, 'from unittest.mock import MagicMock, call, patch\n'), ((9317, 9451), 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', 'PCSMPCInstance.create_instance', ([], {'instance_id': 'test_mpc_id', 'game_name': 'GameNames.LIFT.value', 'mpc_party': 'MPCParty.SERVER', 'num_workers': '(2)'}), '(instance_id=test_mpc_id, game_name=GameNames\n    .LIFT.value, mpc_party=MPCParty.SERVER, num_workers=2)\n', (9347, 9451), False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), ((9874, 9918), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'updated_mpc_instance'}), '(return_value=updated_mpc_instance)\n', (9883, 9918), False, 'from unittest.mock import MagicMock, call, patch\n'), ((10010, 10062), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (10019, 10062), False, 'from unittest.mock import MagicMock, call, patch\n'), ((13281, 13316), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'pl_instance'}), '(return_value=pl_instance)\n', (13290, 13316), False, 'from unittest.mock import MagicMock, call, patch\n'), ((14113, 14148), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'pl_instance'}), '(return_value=pl_instance)\n', (14122, 14148), False, 'from unittest.mock import MagicMock, call, patch\n'), ((15000, 15035), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'pl_instance'}), '(return_value=pl_instance)\n', (15009, 15035), False, 'from unittest.mock import MagicMock, call, patch\n'), ((15908, 15943), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'pl_instance'}), '(return_value=pl_instance)\n', (15917, 15943), False, 'from unittest.mock import MagicMock, call, patch\n'), ((16694, 16729), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'pl_instance'}), '(return_value=pl_instance)\n', (16703, 16729), False, 'from unittest.mock import MagicMock, call, patch\n'), ((17617, 17669), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (17626, 17669), False, 'from unittest.mock import MagicMock, call, patch\n'), ((18387, 18585), 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_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'}), '(instance_id=test_mpc_id, game_name=GameNames\n    .LIFT.value, mpc_party=MPCParty.SERVER, num_workers=self.\n    test_num_containers, status=MPCInstanceStatus.COMPLETED)\n', (18417, 18585), False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), ((18906, 18958), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (18915, 18958), False, 'from unittest.mock import MagicMock, call, patch\n'), ((19020, 19056), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'mpc_instance'}), '(return_value=mpc_instance)\n', (19029, 19056), False, 'from unittest.mock import MagicMock, call, patch\n'), ((20783, 21003), 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_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'}), "(instance_id=test_private_computation_id +\n    '_aggregate_shards', game_name=GameNames.SHARD_AGGREGATOR.value,\n    mpc_party=MPCParty.SERVER, num_workers=2, status=MPCInstanceStatus.FAILED)\n", (20813, 21003), False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), ((21324, 21376), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (21333, 21376), False, 'from unittest.mock import MagicMock, call, patch\n'), ((21438, 21474), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'mpc_instance'}), '(return_value=mpc_instance)\n', (21447, 21474), False, 'from unittest.mock import MagicMock, call, patch\n'), ((22831, 22883), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (22840, 22883), False, 'from unittest.mock import MagicMock, call, patch\n'), ((24558, 24569), 'unittest.mock.MagicMock', 'MagicMock', ([], {}), '()\n', (24567, 24569), False, 'from unittest.mock import MagicMock, call, patch\n'), ((24614, 24625), 'libfb.py.asyncio.mock.AsyncMock', 'AsyncMock', ([], {}), '()\n', (24623, 24625), False, 'from libfb.py.asyncio.mock import AsyncMock\n'), ((26966, 27151), 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_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'}), "(instance_id='test_mpc_id', game_name=\n    GameNames.SHARD_AGGREGATOR.value, mpc_party=MPCParty.SERVER,\n    num_workers=2, status=MPCInstanceStatus.FAILED)\n", (26996, 27151), False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), ((27467, 27710), 'fbpcs.pid.entity.pid_instance.PIDInstance', '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'}), "(instance_id='test_pid_id', protocol=PIDProtocol.UNION_PID,\n    pid_role=PIDRole.PUBLISHER, num_shards=4, input_path='input',\n    output_path='output', stages_containers={}, stages_status={}, status=\n    PIDInstanceStatus.COMPLETED)\n", (27478, 27710), False, 'from fbpcs.pid.entity.pid_instance import PIDInstance, PIDInstanceStatus, PIDProtocol, PIDRole\n'), ((28245, 28297), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (28254, 28297), False, 'from unittest.mock import MagicMock, call, patch\n'), ((29865, 29917), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (29874, 29917), False, 'from unittest.mock import MagicMock, call, patch\n'), ((30893, 30904), 'unittest.mock.MagicMock', 'MagicMock', ([], {}), '()\n', (30902, 30904), False, 'from unittest.mock import MagicMock, call, patch\n'), ((32731, 32920), 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', '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'}), '(instance_id=test_mpc_id, game_name=\n    test_game_name, mpc_party=test_mpc_party, num_workers=self.\n    test_num_containers, status=MPCInstanceStatus.STARTED)\n', (32761, 32920), False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), ((33296, 33348), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'private_computation_instance'}), '(return_value=private_computation_instance)\n', (33305, 33348), False, 'from unittest.mock import MagicMock, call, patch\n'), ((33488, 33678), 'fbpcs.common.entity.pcs_mpc_instance.PCSMPCInstance.create_instance', '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'}), '(instance_id=test_mpc_id, game_name=\n    test_game_name, mpc_party=test_mpc_party, num_workers=self.\n    test_num_containers, status=MPCInstanceStatus.CANCELED)\n', (33518, 33678), False, 'from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance\n'), ((33805, 33850), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'mpc_instance_canceled'}), '(return_value=mpc_instance_canceled)\n', (33814, 33850), False, 'from unittest.mock import MagicMock, call, patch\n'), ((33949, 33994), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'mpc_instance_canceled'}), '(return_value=mpc_instance_canceled)\n', (33958, 33994), False, 'from unittest.mock import MagicMock, call, patch\n'), ((35627, 35683), 'fbpcs.private_computation.service.utils.gen_mpc_game_args_to_retry', 'gen_mpc_game_args_to_retry', (['private_computation_instance'], {}), '(private_computation_instance)\n', (35653, 35683), 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'), ((36131, 36672), 'fbpcs.private_computation.entity.private_computation_instance.PrivateComputationInstance', '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'}), '(instance_id=self.test_private_computation_id,\n    role=role, instances=instances or [], status=status, status_update_ts=\n    1600000000, num_pid_containers=self.test_num_containers,\n    num_mpc_containers=self.test_num_containers, concurrency=self.\n    test_concurrency, num_files_per_mpc_container=NUM_NEW_SHARDS_PER_FILE,\n    game_type=PrivateComputationGameType.LIFT, input_path=self.\n    test_input_path, output_dir=self.test_output_dir, fail_fast=True,\n    k_anonymity_threshold=DEFAULT_K_ANONYMITY_THRESHOLD)\n', (36157, 36672), False, 'from fbpcs.private_computation.entity.private_computation_instance import PrivateComputationGameType, PrivateComputationInstance, PrivateComputationInstanceStatus, PrivateComputationRole, UnionedPCInstance\n'), ((12435, 12470), 'unittest.mock.MagicMock', 'MagicMock', ([], {'return_value': 'pl_instance'}), '(return_value=pl_instance)\n', (12444, 12470), False, 'from unittest.mock import MagicMock, call, patch\n'), ((25379, 25673), 'fbpcs.private_computation.service.utils.create_and_start_mpc_instance', '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'}), '(mpc_svc=mock_mpc_svc, instance_id=instance_id,\n    game_name=game_name, mpc_party=mpc_party, num_containers=num_containers,\n    binary_version=binary_version, container_timeout=\n    DEFAULT_CONTAINER_TIMEOUT_IN_SEC, server_ips=server_ips, game_args=\n    game_args)\n', (25408, 25673), 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'), ((25903, 26027), 'unittest.mock.call', 'call', ([], {'instance_id': 'instance_id', 'game_name': 'game_name', 'mpc_party': 'mpc_party', 'num_workers': 'num_containers', 'game_args': 'game_args'}), '(instance_id=instance_id, game_name=game_name, mpc_party=mpc_party,\n    num_workers=num_containers, game_args=game_args)\n', (25907, 26027), False, 'from unittest.mock import MagicMock, call, patch\n'), ((26221, 26344), 'unittest.mock.call', 'call', ([], {'instance_id': 'instance_id', 'server_ips': 'server_ips', 'timeout': 'DEFAULT_CONTAINER_TIMEOUT_IN_SEC', 'version': 'binary_version'}), '(instance_id=instance_id, server_ips=server_ips, timeout=\n    DEFAULT_CONTAINER_TIMEOUT_IN_SEC, version=binary_version)\n', (26225, 26344), False, 'from unittest.mock import MagicMock, call, patch\n'), ((26617, 26692), 'fbpcs.private_computation.service.utils.map_private_computation_role_to_mpc_party', 'map_private_computation_role_to_mpc_party', (['PrivateComputationRole.PUBLISHER'], {}), '(PrivateComputationRole.PUBLISHER)\n', (26658, 26692), 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'), ((26771, 26844), 'fbpcs.private_computation.service.utils.map_private_computation_role_to_mpc_party', 'map_private_computation_role_to_mpc_party', (['PrivateComputationRole.PARTNER'], {}), '(PrivateComputationRole.PARTNER)\n', (26812, 26844), 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'), ((28334, 28407), 'unittest.mock.patch.object', 'patch.object', (['CppLiftIdSpineCombinerService', '"""combine_on_container_async"""'], {}), "(CppLiftIdSpineCombinerService, 'combine_on_container_async')\n", (28346, 28407), False, 'from unittest.mock import MagicMock, call, patch\n'), ((28460, 28520), 'unittest.mock.patch.object', 'patch.object', (['CppShardingService', '"""shard_on_container_async"""'], {}), "(CppShardingService, 'shard_on_container_async')\n", (28472, 28520), False, 'from unittest.mock import MagicMock, call, patch\n'), ((29954, 30027), 'unittest.mock.patch.object', 'patch.object', (['CppLiftIdSpineCombinerService', '"""combine_on_container_async"""'], {}), "(CppLiftIdSpineCombinerService, 'combine_on_container_async')\n", (29966, 30027), False, 'from unittest.mock import MagicMock, call, patch\n'), ((30080, 30140), 'unittest.mock.patch.object', 'patch.object', (['CppShardingService', '"""shard_on_container_async"""'], {}), "(CppShardingService, 'shard_on_container_async')\n", (30092, 30140), False, 'from unittest.mock import MagicMock, call, patch\n'), ((4415, 4504), 'fbpcs.onedocker_binary_config.OneDockerBinaryConfig', 'OneDockerBinaryConfig', ([], {'tmp_directory': '"""/test_tmp_directory/"""', 'binary_version': '"""latest"""'}), "(tmp_directory='/test_tmp_directory/', binary_version=\n    'latest')\n", (4436, 4504), False, 'from fbpcs.onedocker_binary_config import OneDockerBinaryConfig\n'), ((11468, 11539), 'libfb.py.asyncio.mock.AsyncMock', 'AsyncMock', ([], {'side_effect': '(lambda pc_instance, *args, **kwargs: pc_instance)'}), '(side_effect=lambda pc_instance, *args, **kwargs: pc_instance)\n', (11477, 11539), False, 'from libfb.py.asyncio.mock import AsyncMock\n'), ((34846, 34943), 'fbpcp.entity.container_instance.ContainerInstance', 'ContainerInstance', ([], {'instance_id': '"""container_instance_0"""', 'status': 'ContainerInstanceStatus.FAILED'}), "(instance_id='container_instance_0', status=\n    ContainerInstanceStatus.FAILED)\n", (34863, 34943), False, 'from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus\n'), ((35015, 35115), 'fbpcp.entity.container_instance.ContainerInstance', 'ContainerInstance', ([], {'instance_id': '"""container_instance_1"""', 'status': 'ContainerInstanceStatus.COMPLETED'}), "(instance_id='container_instance_1', status=\n    ContainerInstanceStatus.COMPLETED)\n", (35032, 35115), False, 'from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus\n')]