def setUp(self, onedocker_service) -> None: self.onedocker_service = onedocker_service self.test_num_containers = 2 self.onedocker_binary_config_map = defaultdict( lambda: OneDockerBinaryConfig( tmp_directory="/test_tmp_directory/", binary_version="latest", repository_path="test_path/", )) self.stage_svc = PrepareDataStageService( self.onedocker_service, self.onedocker_binary_config_map)
def setUp(self, mock_mpc_svc: MPCService) -> None: self.mock_mpc_svc = mock_mpc_svc self.mock_mpc_svc.create_instance = MagicMock() onedocker_binary_config_map = defaultdict( lambda: OneDockerBinaryConfig( tmp_directory="/test_tmp_directory/", binary_version="latest", repository_path="test_path/", )) self.stage_svc = ComputeMetricsStageService( onedocker_binary_config_map, self.mock_mpc_svc)
def setUp(self, mock_onedocker_svc, mock_storage_svc) -> None: self.mock_onedocker_svc = mock_onedocker_svc self.mock_storage_svc = mock_storage_svc self.onedocker_binary_config = OneDockerBinaryConfig( tmp_directory="/tmp", binary_version="latest", repository_path="test_path/", ) self.binary_name = "data_processing/pid_preparer" self.onedocker_binary_config_map = { self.binary_name: self.onedocker_binary_config } self.input_path = "in" self.output_path = "out" self.pc_instance_id = "test_instance_123" self.container_timeout = 43200
async def _run_sub_test( wait_for_containers: bool, expected_container_status: ContainerInstanceStatus, ) -> None: with patch.object(PIDStage, "update_instance_containers"): test_onedocker_binary_config = OneDockerBinaryConfig( tmp_directory="/test_tmp_directory/", binary_version="latest", repository_path="test_path/", ) container = ContainerInstance( instance_id="123", ip_address="192.0.2.0", status=expected_container_status, ) mock_sharder.return_value = [container] stage = PIDShardStage( stage=UnionPIDStage.PUBLISHER_SHARD, instance_repository=mock_instance_repo, storage_svc=mock_storage_svc, onedocker_svc=mock_onedocker_svc, onedocker_binary_config=test_onedocker_binary_config, ) test_input_path = "foo" test_output_path = "bar" test_num_shards = 1 test_hmac_key = "CoXbp7BOEvAN9L1CB2DAORHHr3hB7wE7tpxMYm07tc0=" shard_path = PIDShardStage.get_sharded_filepath( test_output_path, 0) self.assertEqual(f"{test_output_path}_0", shard_path) res = await stage.shard( "123", test_input_path, test_output_path, test_num_shards, test_hmac_key, wait_for_containers=wait_for_containers, ) self.assertEqual( PIDStage.get_stage_status_from_containers([container]), res, ) mock_sharder.assert_called_once()
def setUp(self, mock_onedocker_service, mock_storage_service) -> None: self.mock_onedocker_svc = mock_onedocker_service self.mock_storage_svc = mock_storage_service self.test_num_containers = 1 self.onedocker_binary_config_map = defaultdict( lambda: OneDockerBinaryConfig( tmp_directory="/test_tmp_directory/", binary_version="latest", repository_path="test_path/", ) ) self.server_ips = [f"192.0.2.{i}" for i in range(self.test_num_containers)] self.input_path = "in" self.output_path = "out" self.pc_instance_id = "test_instance_123" self.port = 15200 self.use_row_numbers = True
def setUp(self, mock_onedocker_svc, mock_storage_svc) -> None: self.mock_onedocker_svc = mock_onedocker_svc self.mock_storage_svc = mock_storage_svc self.onedocker_binary_config = OneDockerBinaryConfig( tmp_directory="/tmp", binary_version="latest", repository_path="test_path/", ) self.onedocker_binary_config = self.onedocker_binary_config self.binary_name = "data_processing/sharder_hashed_for_pid" self.onedocker_binary_config_map = { self.binary_name: self.onedocker_binary_config } self.input_path = "in" self.output_path = "out" self.pc_instance_id = "test_instance_123" self.container_timeout = 789 self.test_hmac_key = "CoXbp7BOEvAN9L1CB2DAORHHr3hB7wE7tpxMYm07tc0="
def setUp(self) -> None: self._pc_instance = PrivateComputationInstance( instance_id="123", role=PrivateComputationRole.PARTNER, instances=[], status=PrivateComputationInstanceStatus. INPUT_DATA_VALIDATION_STARTED, status_update_ts=1600000000, num_pid_containers=1, num_mpc_containers=1, num_files_per_mpc_container=1, game_type=PrivateComputationGameType.LIFT, input_path= "https://a-test-bucket.s3.us-west-2.amazonaws.com/lift/test/input_data1.csv", output_dir="789", ) self.onedocker_binary_config_map = defaultdict( lambda: OneDockerBinaryConfig( tmp_directory="/test_tmp_directory/", binary_version="latest", repository_path="test_path/", ))
async def _run_sub_test(wait_for_containers: bool, ) -> None: ip = "192.0.2.0" container = ContainerInstance(instance_id="123", ip_address=ip) mock_onedocker_svc.start_containers = MagicMock( return_value=[container]) mock_onedocker_svc.wait_for_pending_containers = AsyncMock( return_value=[container]) container.status = (ContainerInstanceStatus.COMPLETED if wait_for_containers else ContainerInstanceStatus.STARTED) mock_wait_for_containers_async.return_value = [container] test_onedocker_binary_config = OneDockerBinaryConfig( tmp_directory="/test_tmp_directory/", binary_version="latest", repository_path="test_path/", ) stage = PIDShardStage( stage=UnionPIDStage.PUBLISHER_SHARD, instance_repository=mock_instance_repo, storage_svc=mock_storage_svc, onedocker_svc=mock_onedocker_svc, onedocker_binary_config=test_onedocker_binary_config, ) instance_id = "444" stage_input = PIDStageInput( input_paths=["in"], output_paths=["out"], num_shards=123, instance_id=instance_id, ) # Basic test: All good with patch.object(PIDShardStage, "files_exist") as mock_fe: mock_fe.return_value = True stage = PIDShardStage( stage=UnionPIDStage.PUBLISHER_SHARD, instance_repository=mock_instance_repo, storage_svc=mock_storage_svc, onedocker_svc=mock_onedocker_svc, onedocker_binary_config=test_onedocker_binary_config, ) status = await stage.run( stage_input, wait_for_containers=wait_for_containers) self.assertEqual( PIDStageStatus.COMPLETED if wait_for_containers else PIDStageStatus.STARTED, status, ) mock_onedocker_svc.start_containers.assert_called_once() if wait_for_containers: mock_wait_for_containers_async.assert_called_once() else: mock_wait_for_containers_async.assert_not_called() # instance status is updated to READY, STARTED, then COMPLETED mock_instance_repo.read.assert_called_with(instance_id) self.assertEqual(mock_instance_repo.read.call_count, 4) self.assertEqual(mock_instance_repo.update.call_count, 4) # Input not ready with patch.object(PIDShardStage, "files_exist") as mock_fe: mock_fe.return_value = False status = await stage.run( stage_input, wait_for_containers=wait_for_containers) self.assertEqual(PIDStageStatus.FAILED, status) # Multiple input paths (invariant exception) with patch.object(PIDShardStage, "files_exist") as mock_fe: with self.assertRaises(ValueError): mock_fe.return_value = True stage_input.input_paths = ["in1", "in2"] stage = PIDShardStage( stage=UnionPIDStage.PUBLISHER_SHARD, instance_repository=mock_instance_repo, storage_svc=mock_storage_svc, onedocker_svc=mock_onedocker_svc, onedocker_binary_config=test_onedocker_binary_config, ) status = await stage.run( stage_input, wait_for_containers=wait_for_containers) self.assertEqual( PIDStageStatus.COMPLETED if wait_for_containers else PIDStageStatus.STARTED, status, )
def setUp(self) -> None: self.onedocker_binary_config = OneDockerBinaryConfig( tmp_directory="/tmp/", binary_version="latest", repository_path="test_path/", )
def _build_onedocker_binary_cfg( onedocker_binary_config: Dict[str, Any]) -> OneDockerBinaryConfig: return OneDockerBinaryConfig(**onedocker_binary_config["constructor"])
def setUp(self) -> None: 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", repository_path="test_path/", )) 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.pc_validator_config = PCValidatorConfig(region="us-west-2", ) self.private_computation_service = PrivateComputationService( instance_repository=private_computation_instance_repository, storage_svc=storage_svc, mpc_svc=self.mpc_service, pid_svc=self.pid_service, onedocker_svc=self.onedocker_service, onedocker_binary_config_map=self.onedocker_binary_config_map, pc_validator_config=self.pc_validator_config, ) self.bolt_pcs_client = BoltPCSClient(self.private_computation_service) self.test_instance_id = "test_id" self.test_role = PrivateComputationRole.PUBLISHER self.test_game_type = PrivateComputationGameType.LIFT self.test_input_path = "in_path" self.test_output_path = "out_path" self.test_num_containers = 2 self.test_concurrency = 1 self.test_hmac_key = "CoXbp7BOEvAN9L1CB2DAORHHr3hB7wE7tpxMYm07tc0="