diff --git a/squad_client/commands/download_attachments.py b/squad_client/commands/download_attachments.py new file mode 100644 index 0000000..a09909a --- /dev/null +++ b/squad_client/commands/download_attachments.py @@ -0,0 +1,43 @@ +from squad_client import logging +from squad_client.core.command import SquadClientCommand +from squad_client.core.models import TestRun +from squad_client.shortcuts import download_attachments + + +logger = logging.getLogger(__name__) + + +class DownloadAttachmentsCommand(SquadClientCommand): + command = "download-attachments" + help_text = "download the attachments from a SQUAD testrun" + + def register(self, subparser): + parser = super(DownloadAttachmentsCommand, self).register(subparser) + parser.add_argument( + "--testrun", + help="The SQUAD ID of the testrun", + type=int, + required=True, + ) + parser.add_argument( + "--filenames", + nargs="+", + help="If not all files are required, specify the names of the files to download", + ) + parser.add_argument( + "--debug", + action='store_true', + help="Set debug mode" + ) + + def run(self, args): + if args.debug: + logger.setLevel(logging.DEBUG) + + # Check testrun + testrun = TestRun(args.testrun) + if testrun is None: + logger.error(f"TestRun {args.testrun} not found") + + # Check if requested files exist + return download_attachments(testrun, args.filenames) diff --git a/squad_client/shortcuts.py b/squad_client/shortcuts.py index 7ad1b81..6d39866 100644 --- a/squad_client/shortcuts.py +++ b/squad_client/shortcuts.py @@ -315,6 +315,24 @@ def get_build(build_version, project): return first(project.builds(**filters)) +def download_attachments(testrun, filenames=None): + attachment_filenames = [attachment.filename for attachment in testrun.attachments] + if filenames: + for filename in filenames: + if filename not in attachment_filenames: + logger.error(f"Error: Attachment filename {filename} not found in TestRun {testrun.id}. " + f"Attachments in TestRun are: {attachment_filenames}") + return False + + for attachment in testrun.attachments: + if not filenames or attachment.filename in filenames: + logger.debug(f"Downloading {attachment.filename}") + with open(attachment.filename, 'wb') as fp: + fp.write(attachment.read()) + + return True + + def download_tests(project, build, filter_envs=None, filter_suites=None, format_string=None, output_filename=None): all_environments = project.environments(count=ALL) all_suites = project.suites(count=ALL) diff --git a/tests/test_shortcuts.py b/tests/test_shortcuts.py index 944caf3..6e0790c 100644 --- a/tests/test_shortcuts.py +++ b/tests/test_shortcuts.py @@ -1,4 +1,5 @@ import logging +import os from unittest import TestCase @@ -13,6 +14,7 @@ submit_job, create_or_update_project, watchjob, + download_attachments, download_tests, register_callback, ) @@ -333,6 +335,43 @@ def test_no_overwrite(self): project.delete() +class DownloadAttachmentsShortcutTest(TestCase): + def setUp(self): + self.squad = Squad() + SquadApi.configure(url="http://localhost:%s" % settings.DEFAULT_SQUAD_PORT) + group = self.squad.group("my_group") + project = group.project("my_project") + build = project.build("my_build") + self.testrun = build.testruns()[3] + + def tearDown(self): + for filename in ["foo1.txt", "foo2.txt"]: + if os.path.isfile(filename): + os.remove(filename) + + def test_download_all(self): + expected_files = ["foo1.txt", "foo2.txt"] + success = download_attachments(self.testrun) + self.assertTrue(success) + # Check all files downloaded + for filename in expected_files: + self.assertTrue(os.path.isfile(filename)) + + def test_download_file_list(self): + filenames = ["foo2.txt"] + success = download_attachments(self.testrun, filenames) + self.assertTrue(success) + # Check foo1.txt wasn't downloaded but foo2.txt was + self.assertFalse(os.path.isfile("foo1.txt")) + self.assertTrue(os.path.isfile("foo2.txt")) + + def test_download_file_missing(self): + filenames = ["foo1.txt", "foo3.txt"] + success = download_attachments(self.testrun, filenames) + # Download should fail + self.assertFalse(success) + + class DownloadTestsShortcutTest(TestCase): def setUp(self): self.squad = Squad()