def test_create_job_prompts(runner, mock_create, patch_os_getcwd):
cli_inputs = ["job", "create"]
prompt_inputs = ["test-job", "test-gcp-project"]
inputs = "\n".join(prompt_inputs)
result = runner.invoke(cli.main, cli_inputs, input=inputs)
core_testing.assert_execution_success(result)
exp_output = (
"Name of your new job: {0}\n"
"Name of the GCP project the job should be created in: "
"{1}\n".format(*prompt_inputs)
)
assert exp_output == result.output
known_kwargs = {
"job_name": "test-job",
"gcp_project": "test-gcp-project",
"use_defaults": False,
}
mock_create.assert_called_once_with((), known_kwargs, patch_os_getcwd)
def test_profile_memory_per_line(runner, mocker, minimal_mock_klio_config):
mock_profile = mocker.patch.object(cli, "_profile")
result = runner.invoke(cli.profile_memory_per_line, [])
core_testing.assert_execution_success(result)
exp_kwargs = {
"get_maximum": False,
"per_element": False,
"image_tag": None,
"force_build": False,
"input_file": None,
"output_file": None,
"show_logs": False,
"entity_ids": (),
}
mock_profile.assert_called_once_with("memory-per-line",
minimal_mock_klio_config.klio_config,
minimal_mock_klio_config.meta,
**exp_kwargs)
def test_create_job(runner, mock_create, patch_os_getcwd):
cli_inputs = [
"job",
"create",
"--job-name",
"test-job",
"--gcp-project",
"test-gcp-project",
"--use-defaults",
]
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
known_kwargs = {
"job_name": "test-job",
"gcp_project": "test-gcp-project",
"use_defaults": True,
}
mock_create.assert_called_once_with((), known_kwargs, patch_os_getcwd)
def test_stop_job(
config_override,
runner,
mock_stop,
mocker,
config_file,
mock_get_config_job_dir,
patch_os_getcwd,
pipeline_config_dict,
mock_klio_config,
):
config_data = {
"job_name": "test-job",
"version": 1,
"pipeline_options": pipeline_config_dict,
"job_config": {
"inputs": [{
"topic": "foo-topic",
"subscription": "foo-sub",
"data_location": "foo-input-location",
}],
"outputs": [{
"topic": "foo-topic-output",
"data_location": "foo-output-location",
}],
},
}
mock_klio_config.setup(config_data, config_file, config_override)
cli_inputs = ["job", "stop"]
if config_override:
cli_inputs.extend(["--config-file", config_override])
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
mock_klio_config.assert_calls()
mock_stop.assert_called_once_with("test-job", "test-project",
"us-central1", "cancel")
def test_run_pipeline(
image_tag,
direct_runner,
update,
blocking,
cli_runner,
mock_klio_config,
patch_run_basic_pipeline,
):
mock_compare_runtime_to_buildtime_config.return_value = True
runtime_conf = cli.RuntimeConfig(image_tag=None,
direct_runner=False,
update=None,
blocking=None)
cli_inputs = []
if image_tag:
cli_inputs.extend(["--image-tag", image_tag])
runtime_conf = runtime_conf._replace(image_tag=image_tag)
if direct_runner:
cli_inputs.append("--direct-runner")
runtime_conf = runtime_conf._replace(direct_runner=True)
if update:
cli_inputs.append("--update")
runtime_conf = runtime_conf._replace(update=True)
if update is False:
cli_inputs.append("--no-update")
if not update: # if none or false
runtime_conf = runtime_conf._replace(update=False)
if blocking:
cli_inputs.append("--blocking")
runtime_conf._replace(blocking=True)
if blocking is False:
cli_inputs.append("--no-blocking")
if not blocking:
runtime_conf._replace(blocking=False)
result = cli_runner.invoke(cli.run_pipeline, cli_inputs)
core_testing.assert_execution_success(result)
mock_klio_config.assert_calls()
patch_run_basic_pipeline.assert_called_once_with()
def test_profile_cpu(runner, mocker, minimal_mock_klio_config):
mock_profile = mocker.patch.object(cli, "_profile")
result = runner.invoke(cli.profile_cpu, [])
core_testing.assert_execution_success(result)
exp_kwargs = {
"interval": 0.1,
"plot_graph": False,
"image_tag": None,
"force_build": False,
"input_file": None,
"output_file": None,
"show_logs": False,
"entity_ids": (),
}
mock_profile.assert_called_once_with("cpu",
minimal_mock_klio_config.klio_config,
minimal_mock_klio_config.meta,
**exp_kwargs)
def test_test_job_conf_override(
config_file_override,
pytest_args,
mocker,
monkeypatch,
config,
cli_runner,
patch_klio_config,
tmpdir,
):
mock_get_config = mocker.Mock()
monkeypatch.setattr(cli, "_get_config", mock_get_config)
mock_get_config.return_value = config
mock_test = mocker.Mock()
monkeypatch.setattr(pytest, "main", mock_test)
mock_test.return_value = 0
cli_inputs = []
temp_dir = tmpdir.mkdir("testing")
temp_dir_str = str(temp_dir)
monkeypatch.setattr(os, "getcwd", lambda: temp_dir_str)
config_file = "klio-job.yaml"
if config_file_override:
# create a tmp file else click will complain it doesn't exist
config_file = _create_tmp_config_file(temp_dir_str,
config_file_override, config)
cli_inputs.extend(["--config-file", config_file])
cli_inputs.extend(pytest_args)
result = cli_runner.invoke(cli.test_job, cli_inputs)
core_testing.assert_execution_success(result)
assert "true" == os.environ["KLIO_TEST_MODE"]
mock_test.assert_called_once_with(pytest_args)
mock_get_config.assert_called_once_with(config_file)
def test_delete_job(
config_override,
mocker,
monkeypatch,
runner,
config_file,
mock_delete,
patch_os_getcwd,
mock_klio_config,
):
config = {
"job_name": "test-job",
"version": 1,
"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_obj = mock_klio_config.setup(config, config_file, config_override)
cli_inputs = ["job", "delete"]
if config_override:
cli_inputs.extend(["--config-file", config_override])
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
mock_klio_config.assert_calls()
mock_delete.assert_called_once_with(conf_obj)
mock_delete.return_value.delete.assert_called_once_with()
def test_run_pipeline_conf_override(
config_file_override,
cli_runner,
config,
mock_klio_config,
patch_run_basic_pipeline,
caplog,
tmpdir,
monkeypatch,
):
cli_inputs = []
temp_dir = tmpdir.mkdir("testing123")
temp_dir_str = str(temp_dir)
monkeypatch.setattr(os, "getcwd", lambda: temp_dir_str)
exp_conf_file = "klio-job.yaml"
if config_file_override:
exp_conf_file = os.path.join(temp_dir_str, config_file_override)
cli_inputs.extend(["--config-file", exp_conf_file])
# create a tmp file else click will complain it doesn't exist
with open(exp_conf_file, "w") as f:
yaml.dump(config, f)
mock_klio_config.setup(_config(), "klio-job.yaml", exp_conf_file)
else:
mock_klio_config.setup(_config(), "klio-job.yaml")
result = cli_runner.invoke(cli.run_pipeline, cli_inputs)
core_testing.assert_execution_success(result)
mock_klio_config.assert_calls()
patch_run_basic_pipeline.assert_called_once_with()
assert 0 == len(caplog.records)
def test_test_job(
runner,
mocker,
config_file,
patch_os_getcwd,
pytest_args,
conf_override,
image_tag,
mock_get_git_sha,
mock_warn_if_py2_job,
mock_get_config_job_dir,
):
mock_test_pipeline = mocker.patch.object(cli.job_commands.test,
"TestPipeline")
mock_test_pipeline.return_value.run.return_value = 0
mock_get_config_job_dir.return_value = (
patch_os_getcwd,
conf_override or config_file,
)
cli_inputs = ["job", "test"]
if image_tag:
cli_inputs.extend(["--image-tag", image_tag])
if conf_override:
cli_inputs.extend(["--config-file", conf_override])
cli_inputs.extend(pytest_args)
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",
}],
},
}
mock_get_config = mocker.patch.object(core_utils, "get_config_by_path")
# deepcopy since KlioConfig will pop keys
mock_get_config.return_value = config_data
conf = kconfig.KlioConfig(copy.deepcopy(config_data))
mock_klio_config = mocker.patch.object(core_utils.config, "KlioConfig")
mock_klio_config.return_value = conf
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
exp_image_tag = image_tag or mock_get_git_sha.return_value
if conf_override:
exp_image_tag = "{}-{}".format(exp_image_tag, conf_override)
mock_get_config_job_dir.assert_called_once_with(None, conf_override)
mock_warn_if_py2_job.assert_called_once_with(patch_os_getcwd)
if not image_tag:
mock_get_git_sha.assert_called_once_with(patch_os_getcwd)
else:
mock_get_git_sha.assert_not_called()
mock_get_config.assert_called_once_with(conf_override or config_file)
mock_klio_config.assert_called_once_with(config_data)
exp_docker_runtime_config = cli.DockerRuntimeConfig(
image_tag=exp_image_tag,
force_build=False,
config_file_override=conf_override,
)
mock_test_pipeline.assert_called_once_with(patch_os_getcwd, conf,
exp_docker_runtime_config)
mock_test_pipeline.return_value.run.assert_called_once_with(
pytest_args=pytest_args)
def test_publish(
force,
ping,
top_down,
bottom_up,
non_klio,
conf_override,
runner,
mocker,
config_file,
patch_os_getcwd,
mock_klio_config,
):
mock_publish = mocker.patch.object(cli.message_commands.publish,
"publish_messages")
config = {
"job_name": "test-job",
"pipeline_options": {
"worker_harness_container_image": "gcr.register.io/squad/feature"
},
"job_config": {
"inputs": [{
"topic": "foo-topic",
"subscription": "foo-sub",
"data_location": "foo-input-location",
}],
"outputs": [{
"topic": "foo-topic-output",
"data_location": "foo-output-location",
}],
},
}
cli_inputs = ["message", "publish", "deadb33f"]
if force:
cli_inputs.append("--force")
if ping:
cli_inputs.append("--ping")
if top_down:
cli_inputs.append("--top-down")
if bottom_up:
cli_inputs.append("--bottom-up")
if non_klio:
cli_inputs.append("--non-klio")
config["job_config"]["allow_non_klio_messages"] = True
if conf_override:
cli_inputs.extend(["--config-file", conf_override])
conf = mock_klio_config.setup(config, config_file)
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
mock_publish.assert_called_once_with(
conf,
("deadb33f", ),
force,
ping,
top_down,
non_klio,
)
def test_run_job(
runner,
mocker,
tmpdir,
is_gcp,
direct_runner,
image_tag,
config_override,
config_file,
mock_get_git_sha,
mock_klio_config,
is_job_dir_override,
):
mock_run = mocker.patch.object(cli.job_commands.run.RunPipeline, "run")
mock_run.return_value = 0
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",
}],
},
}
job_dir_override = None
if is_job_dir_override:
temp_dir = tmpdir.mkdir("testing12345")
job_dir_override = str(temp_dir)
mock_klio_config.setup(config_data, config_file, config_override,
job_dir_override)
cli_inputs = ["job", "run"]
if image_tag:
cli_inputs.extend(["--image-tag", image_tag])
if direct_runner:
cli_inputs.append("--direct-runner")
if config_override:
cli_inputs.extend(["--config-file", config_override])
if job_dir_override:
cli_inputs.extend(["--job-dir", job_dir_override])
exp_image_tag = image_tag or mock_get_git_sha.return_value
if config_override:
exp_image_tag = "{}-{}".format(exp_image_tag, config_override)
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
mock_klio_config.assert_calls()
mock_run.assert_called_once_with()
mock_get_git_sha.assert_called_once_with(mock_klio_config.meta.job_dir,
image_tag)
def test_deploy_job(
runner,
mocker,
is_gcp,
direct_runner,
config_override,
image_tag,
mock_stop,
config_file,
patch_os_getcwd,
pipeline_config_dict,
mock_get_git_sha,
mock_klio_config,
):
mock_run = mocker.patch.object(cli.job_commands.run.RunPipeline, "run")
mock_run.return_value = 0
config_data = {
"job_name": "test-job",
"version": 1,
"pipeline_options": pipeline_config_dict,
"job_config": {
"inputs": [
{
"topic": "foo-topic",
"subscription": "foo-sub",
"data_location": "foo-input-location",
}
],
"outputs": [
{
"topic": "foo-topic-output",
"data_location": "foo-output-location",
}
],
},
}
mock_klio_config.setup(
copy.deepcopy(config_data), config_file, config_override
)
cli_inputs = ["job", "deploy"]
if image_tag:
cli_inputs.extend(["--image_tag", image_tag])
if direct_runner:
cli_inputs.append("--direct-runner")
if config_override:
cli_inputs.extend(["--config-file", config_override])
exp_image_tag = image_tag or mock_get_git_sha.return_value
if config_override:
exp_image_tag = "{}-{}".format(exp_image_tag, config_override)
result = runner.invoke(cli.main, cli_inputs)
core_testing.assert_execution_success(result)
assert "" == result.output
mock_get_git_sha.assert_called_once_with(patch_os_getcwd, image_tag)
mock_klio_config.assert_calls()
mock_stop.assert_called_once_with(
"test-job", "test-project", "us-central1", "cancel"
)
mock_run.assert_called_once_with()
