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()
