def klio_pipeline(mock_os_environ, config, mocker, monkeypatch):
job_dir = "/test/dir/jobs/test_run_job"
runtime_config = cli.DockerRuntimeConfig(
image_tag="foo-1234",
force_build=False,
config_file_override=None,
)
profile_config = cli.ProfileConfig(
input_file="input.txt",
output_file="output.txt",
show_logs=False,
entity_ids=(),
)
return profile.ProfilePipeline(job_dir, config, runtime_config,
profile_config)
def test_profile(
image_tag,
config_override,
config_file,
patch_os_getcwd,
pipeline_config_dict,
mock_get_git_sha,
mocker,
mock_klio_config,
):
mock_req_profile_input = mocker.patch.object(
cli, "_require_profile_input_data")
config_data = {
"job_name": "test-job",
"pipeline_options": {
"worker_harness_container_image": "gcr.register.io/squad/feature",
"project": "test-project",
"region": "boonies",
"staging_location": "gs://somewhere/over/the/rainbow",
"temp_location": "gs://somewhere/over/the/rainbow",
},
"job_config": {
"inputs": [{
"topic": "foo-topic",
"subscription": "foo-sub",
"data_location": "foo-input-location",
}],
"outputs": [{
"topic": "foo-topic-output",
"data_location": "foo-output-location",
}],
},
}
conf = mock_klio_config.setup(config_data, config_file, config_override)
mock_pipeline = mocker.patch.object(cli.job_commands.profile,
"ProfilePipeline")
exp_image_tag = image_tag or mock_get_git_sha.return_value
if config_override:
exp_image_tag = "{}-{}".format(exp_image_tag, config_override)
exp_basic_config = cli.DockerRuntimeConfig(
image_tag=exp_image_tag,
force_build=False,
config_file_override=config_override,
)
exp_prof_config = cli.ProfileConfig(
input_file="input.txt",
output_file="output.txt",
show_logs=False,
entity_ids=(),
)
kwargs = {
"input_file": "input.txt",
"output_file": "output.txt",
"entity_ids": (),
"force_build": False,
"show_logs": False,
"entity_ids": (),
}
cli._profile(subcommand="foo",
klio_config=conf,
image_tag=image_tag,
config_meta=mock_klio_config.meta,
**kwargs)
mock_req_profile_input.assert_called_once_with("input.txt", ())
exp_call = mocker.call(
mock_klio_config.meta.job_dir,
conf,
exp_basic_config,
exp_prof_config,
)
assert 1 == mock_pipeline.call_count
assert exp_call == mock_pipeline.call_args
mock_pipeline.return_value.run.assert_called_once_with(what="foo",
subcommand_flags={})
def test_get_command(
klio_pipeline,
input_file,
exp_input_file_flags,
output_file,
exp_output_file_flags,
show_logs,
exp_show_logs_flags,
entity_ids,
exp_entity_id_args,
subcommand,
mocker,
monkeypatch,
):
mock_parse_cpu_flags = mocker.Mock()
mock_parse_cpu_flags.return_value = ["cpu"]
monkeypatch.setattr(klio_pipeline, "_parse_cpu_flags",
mock_parse_cpu_flags)
mock_parse_memory_flags = mocker.Mock()
mock_parse_memory_flags.return_value = ["memory"]
monkeypatch.setattr(klio_pipeline, "_parse_memory_flags",
mock_parse_memory_flags)
mock_parse_memory_per_line_flags = mocker.Mock()
mock_parse_memory_per_line_flags.return_value = ["memory-per-line"]
monkeypatch.setattr(
klio_pipeline,
"_parse_memory_per_line_flags",
mock_parse_memory_per_line_flags,
)
mock_parse_timeit_flags = mocker.Mock()
mock_parse_timeit_flags.return_value = ["timeit"]
monkeypatch.setattr(klio_pipeline, "_parse_timeit_flags",
mock_parse_timeit_flags)
profile_conf = cli.ProfileConfig(
input_file=input_file,
output_file=output_file,
entity_ids=entity_ids,
show_logs=show_logs,
)
monkeypatch.setattr(klio_pipeline, "profile_config", profile_conf)
exp_command = ["profile"]
if subcommand:
exp_command.append(subcommand)
exp_command.extend(exp_input_file_flags)
exp_command.extend(exp_output_file_flags)
exp_command.extend(exp_show_logs_flags)
exp_command.extend(exp_entity_id_args)
actual_command = klio_pipeline._get_command(subcommand, {})
assert sorted(exp_command) == sorted(actual_command)
if subcommand == "cpu":
mock_parse_cpu_flags.assert_called_once_with({})
mock_parse_memory_flags.assert_not_called()
mock_parse_memory_per_line_flags.assert_not_called()
mock_parse_timeit_flags.assert_not_called()
elif subcommand == "memory":
mock_parse_cpu_flags.assert_not_called()
mock_parse_memory_flags.assert_called_once_with({})
mock_parse_memory_per_line_flags.assert_not_called()
mock_parse_timeit_flags.assert_not_called()
elif subcommand == "memory-per-line":
mock_parse_cpu_flags.assert_not_called()
mock_parse_memory_flags.assert_not_called()
mock_parse_memory_per_line_flags.assert_called_once_with({})
mock_parse_timeit_flags.assert_not_called()
elif subcommand == "timeit":
mock_parse_cpu_flags.assert_not_called()
mock_parse_memory_flags.assert_not_called()
mock_parse_memory_per_line_flags.assert_not_called()
mock_parse_timeit_flags.assert_called_once_with({})
else:
mock_parse_cpu_flags.assert_not_called()
mock_parse_memory_flags.assert_not_called()
mock_parse_memory_per_line_flags.assert_not_called()
mock_parse_timeit_flags.assert_not_called()