def test_print_log_urls(self, print_log_urls_mock) -> None: argv = [ "print_log_urls", "instance123", f"--config={self.temp_filename}", ] pc_cli.main(argv) print_log_urls_mock.assert_called_once()
def test_get_mpc(self, get_mpc_mock) -> None: argv = [ "get_mpc", "instance123", f"--config={self.temp_filename}", ] pc_cli.main(argv) get_mpc_mock.assert_called_once()
def test_cancel_current_stage(self, cancel_stage_mock) -> None: argv = [ "cancel_current_stage", "instance123", f"--config={self.temp_filename}", ] pc_cli.main(argv) cancel_stage_mock.assert_called_once()
def test_get_attribution_dataset_info( self, get_attribution_dataset_info_mock) -> None: argv = [ "get_attribution_dataset_info", "--dataset_id=dataset123", f"--config={self.temp_filename}", ] pc_cli.main(argv) get_attribution_dataset_info_mock.assert_called_once()
def test_validate(self, validate_mock) -> None: argv = [ "validate", "instance123", f"--config={self.temp_filename}", "--aggregated_result_path=/tmp/aggpath", "--expected_result_path=/tmp/exppath", ] pc_cli.main(argv) validate_mock.assert_called_once()
def test_run_next(self, run_next_mock) -> None: argv = [ "run_next", "instance123", f"--config={self.temp_filename}", ] pc_cli.main(argv) run_next_mock.assert_called_once() run_next_mock.reset_mock() argv.extend([ "--server_ips=192.168.1.1,192.168.1.2", ]) pc_cli.main(argv) run_next_mock.assert_called_once()
def test_run_attribution(self, create_mock) -> None: argv = [ "run_attribution", "--dataset_id=43423422232", "--attribution_rule=last_click_1d", f"--input_path={self.temp_files_paths[0]}", "--aggregation_type=measurement", "--concurrency=4", "--num_files_per_mpc_container=4", f"--config={self.temp_filename}", "--timestamp=1646870400", "--k_anonymity_threshold=0", ] pc_cli.main(argv) create_mock.assert_called_once()
def test_pre_validate_with_minimal_input_paths_args( self, getLoggerMock, pre_validate_service_mock) -> None: getLoggerMock.return_value = getLoggerMock expected_config = ConfigYamlDict.from_file(self.temp_filename) argv = [ "pre_validate", f"--config={self.temp_filename}", f"--input_paths={','.join(self.temp_files_paths)}", ] pc_cli.main(argv) pre_validate_service_mock.pre_validate.assert_called_once_with( config=expected_config, input_paths=self.temp_files_paths, logger=getLoggerMock, )
def test_run_study(self, run_study_mock) -> None: argv = [ "run_study", "12345", f"--config={self.temp_filename}", "--objective_ids=12,34,56,78,90", f"--input_paths={','.join(self.temp_files_paths)},", ] pc_cli.main(argv) run_study_mock.assert_called_once() run_study_mock.reset_mock() argv.extend([ "--tries_per_stage=789", "--dry_run", ]) pc_cli.main(argv) run_study_mock.assert_called_once()
def test_run_instance(self, run_instance_mock) -> None: argv = [ "run_instance", "instance123", f"--config={self.temp_filename}", f"--input_path={self.temp_filename}", "--num_shards=456", ] pc_cli.main(argv) run_instance_mock.assert_called_once() run_instance_mock.reset_mock() argv.extend([ "--tries_per_stage=789", "--dry_run", ]) pc_cli.main(argv) run_instance_mock.assert_called_once()
def test_run_instances(self, run_instances_mock) -> None: # Test with real temporary file and folder argv = [ "run_instances", "instance123,instance456", f"--config={self.temp_filename}", f"--input_paths={','.join(self.temp_files_paths[:2])}", "--num_shards_list=456,789", ] pc_cli.main(argv) run_instances_mock.assert_called_once() run_instances_mock.reset_mock() argv.extend([ "--tries_per_stage=789", "--dry_run", ]) pc_cli.main(argv) run_instances_mock.assert_called_once()
def test_pre_validate_with_pa_args(self, getLoggerMock, pre_validate_service_mock) -> None: getLoggerMock.return_value = getLoggerMock expected_config = ConfigYamlDict.from_file(self.temp_filename) argv = [ "pre_validate", f"--config={self.temp_filename}", "--dataset_id=123", f"--input_path={self.temp_files_paths[0]}", "--timestamp=1651847976", "--attribution_rule=last_click_1d", "--aggregation_type=measurement", "--concurrency=1", "--num_files_per_mpc_container=1", "--k_anonymity_threshold=10", ] pc_cli.main(argv) pre_validate_service_mock.pre_validate.assert_called_once_with( config=expected_config, input_paths=[self.temp_files_paths[0]], logger=getLoggerMock, )
def test_create_instance(self, create_mock) -> None: # Normally such *ultra-specific* test cases against a CLI would be an # antipattern, but since this is our public interface, we want to be # very careful before making that interface change. # Create a temporary folder for testing argv = [ "create_instance", "instance123", f"--config={self.temp_filename}", "--role=PUBLISHER", "--game_type=LIFT", f"--input_path={self.temp_files_paths[0]}", f"--output_dir={self.temp_dir_path}", "--num_pid_containers=111", "--num_mpc_containers=222", ] pc_cli.main(argv) create_mock.assert_called_once() create_mock.reset_mock() argv.extend([ "--attribution_rule=last_click_1d", "--aggregation_type=measurement", "--concurrency=333", "--num_files_per_mpc_container=444", "--padding_size=555", "--k_anonymity_threshold=666", "--hmac_key=bigmac", "--stage_flow=PrivateComputationLocalTestStageFlow", ]) pc_cli.main(argv) create_mock.assert_called_once() # Test with additional input paths of various formats additional_input_paths = [ "https://bucket-name.s3.Region.amazonaws.com/key-name", "https://fbpcs-github-e2e.s3.us-west-2.amazonaws.com/lift/results/partner_expected_result.json", "https://fbpcs-github-e2e.s3.Region.amazonaws.com/lift/results/partner_expected_result.json", "https://s3.Region.amazonaws.com/bucket-name/key-name", "https://fbpcs-github-e2e.s3.us-west-2.amazonaws.com/lift/results/partner_expected_result.json", ] for additional_input_path in additional_input_paths: create_mock.reset_mock() argv = [ "create_instance", "instance123", f"--config={self.temp_filename}", "--role=PUBLISHER", "--game_type=LIFT", f"--input_path={additional_input_path}", f"--output_dir={self.temp_dir_path}", "--num_pid_containers=111", "--num_mpc_containers=222", ] pc_cli.main(argv) create_mock.assert_called_once() create_mock.reset_mock() argv.extend([ "--attribution_rule=last_click_1d", "--aggregation_type=measurement", "--concurrency=333", "--num_files_per_mpc_container=444", "--padding_size=555", "--k_anonymity_threshold=666", "--hmac_key=bigmac", "--stage_flow=PrivateComputationLocalTestStageFlow", ]) pc_cli.main(argv) create_mock.assert_called_once()