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

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

@@ -18,7 +18,7 @@
 
 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()
@@ -29,7 +29,7 @@ def download(self, team: str, name: str, path: str):
         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,6 +38,16 @@ def download(self, team: str, name: str, path: str):
                 )
 
             log.info(f"Downloaded Data Job in {path}/{name}")
+
+            log.info(f"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}")
+
         finally:
             self.__cleanup_archive(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,25 @@ 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_job_key_error(httpserver: PluginHTTPServer, tmpdir: LocalPath):
+    rest_api_url = httpserver.url_for("")
+    temp_dir = tmpdir.mkdir("job")
+
+    httpserver.expect_request(
+        uri="/data-jobs/for-team/test-team/jobs/test-job/sources"
+    ).respond_with_response(Response(status=510))
+
+    httpserver.expect_request(
+        uri="/data-jobs/for-team/test-team/jobs/test-job/sources", method="GET"
+    ).respond_with_response(Response(status=404))
+
+    runner = CliRunner()
+    result = runner.invoke(
+        download_job,
+        ["-n", "test-job", "-t", "test-team", "-u", rest_api_url, "-p", temp_dir],
+    )
+
+    assert result.exit_code != 0
+    assert "what" in result.output and "why" in result.output