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