squad_client/commands/download_attachments: Attachment download command

Add an attachment download command to squad-client to download
attachments from a given TestRun.

Signed-off-by: Katie Worton <[email protected]>

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)

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)

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