vdk-control-cli: download keytab along with downloading job

projects/vdk-control-cli/src/vdk/internal/control/command_groups/job/download_job.py

@@ -18,18 +18,19 @@
 
 class JobDownloadSource:
     def __init__(self, rest_api_url: str):
-        self.sources_api = ApiClientFactory(rest_api_url).get_jobs_sources_api()
+        self.api = ApiClientFactory(rest_api_url)
         self.__job_archive = JobArchive()
 
     @ApiClientErrorDecorator()
     def download(self, team: str, name: str, path: str):
         log.debug(f"Download job {name} of team {team} into parent {path}")
         self.__validate_job_path(path, name)
+        path_for_keytab = path
         job_archive_path = os.path.join(path, f"{name}.zip")
         try:
             log.info(f"Downloading data job {name} in {path}/{name} ...")
             with click_spinner.spinner():
-                response = self.sources_api.data_job_sources_download_with_http_info(
+                response = self.api.get_jobs_sources_api().data_job_sources_download_with_http_info(
                     team_name=team, job_name=name, _preload_content=False
                 )
                 self.__write_response_to_archive(job_archive_path, response)
@@ -38,9 +39,23 @@ def download(self, team: str, name: str, path: str):
                 )
 
             log.info(f"Downloaded Data Job in {path}/{name}")
+            try:
+                self.__download_keytab(team, name, path_for_keytab)
+            except Exception as e:
+                log.warning(f"Failed to download keytab for job {name}. Error: {e}.")
         finally:
             self.__cleanup_archive(job_archive_path)
 
+    def __download_keytab(self, team: str, name: str, path: str):
+        log.info("Downloading keytab...")
+        keytab_file_path = os.path.join(path, f"{name}.keytab")
+        response = self.api.get_jobs_api().data_job_keytab_download_with_http_info(
+            team_name=team, job_name=name, _preload_content=False
+        )
+        with open(keytab_file_path, "wb") as w:
+            w.write(response.raw_data)
+        log.info(f"Saved keytab in {keytab_file_path}")
+
     @staticmethod
     def __write_response_to_archive(job_archive_path: str, response: ApiResponse):
         log.debug(f"Write data job source to {job_archive_path}")

projects/vdk-control-cli/tests/vdk/internal/control/command_groups/job/test_download_job.py

@@ -30,6 +30,11 @@ def test_download_source(httpserver: PluginHTTPServer, tmpdir: LocalPath):
         uri="/data-jobs/for-team/test-team/jobs/test-job/sources", method="GET"
     ).respond_with_data(data)
 
+    key_data = b"key_data"
+    httpserver.expect_request(
+        uri="/data-jobs/for-team/test-team/jobs/test-job/keytab", method="GET"
+    ).respond_with_data(key_data)
+
     runner = CliRunner()
     result = runner.invoke(
         download_job,
@@ -43,6 +48,13 @@ def test_download_source(httpserver: PluginHTTPServer, tmpdir: LocalPath):
     assert pathlib.Path(expected_dir_job).is_dir()
     assert pathlib.Path(expected_dir_job).joinpath("config.ini").is_file()
 
+    # Assert the key was downloaded to the parent directory
+    expected_key_path = os.path.join(temp_dir, "test-job.keytab")
+    assert pathlib.Path(expected_key_path).is_file()
+    with open(expected_key_path, "rb") as key_file:
+        actual_key_data = key_file.read()
+        assert actual_key_data == key_data
+
 
 def test_download_source_dir_exists(httpserver: PluginHTTPServer, tmpdir: LocalPath):
     rest_api_url = httpserver.url_for("")
@@ -70,6 +82,11 @@ def test_download_job_does_not_exist(httpserver: PluginHTTPServer, tmpdir: Local
         uri="/data-jobs/for-team/test-team/jobs/test-job/sources", method="GET"
     ).respond_with_response(Response(status=404))
 
+    # Mock a 404 response for the keytab endpoint
+    httpserver.expect_request(
+        uri="/data-jobs/for-team/test-team/jobs/test-job/keytab", method="GET"
+    ).respond_with_response(Response(status=404))
+
     runner = CliRunner()
     result = runner.invoke(
         download_job,
@@ -82,3 +99,32 @@ def test_download_job_does_not_exist(httpserver: PluginHTTPServer, tmpdir: Local
     assert "The requested resource cannot be found" in result.output
     # assert that vdk does not try cleaning up a non-existent archive
     assert "Cannot cleanup archive" not in result.output
+
+
+def test_download_key_does_error(httpserver: PluginHTTPServer, tmpdir: LocalPath):
+    rest_api_url = httpserver.url_for("")
+    temp_dir = tmpdir.mkdir("foo")
+
+    test_job_zip = test_utils.find_test_resource("job-zip/test-job.zip")
+    data = _read_file(test_job_zip)
+
+    httpserver.expect_request(
+        uri="/data-jobs/for-team/test-team/jobs/test-job/sources", method="GET"
+    ).respond_with_data(data)
+
+    # This is the expected failure for the keytab download
+    httpserver.expect_request(
+        uri="/data-jobs/for-team/test-team/jobs/test-job/keytab", method="GET"
+    ).respond_with_response(Response(status=510))
+
+    runner = CliRunner()
+    result = runner.invoke(
+        download_job,
+        ["-n", "test-job", "-t", "test-team", "-u", rest_api_url, "-p", temp_dir],
+    )
+
+    # Still exiting with 0 since job is downloaded
+    test_utils.assert_click_status(result, 0)
+    # Assert the key was not downloaded to the parent directory
+    expected_key_path = os.path.join(temp_dir, "test-job.keytab")
+    assert pathlib.Path(expected_key_path).exists() is False