From f017ec28f601de9d0e44729638ee6484dac717ac Mon Sep 17 00:00:00 2001
From: Jana Cupova <root@localhost>
Date: Aug 18 2022 09:12:17 +0000
Subject: Fix arch in download-task


Fixes: https://pagure.io/koji/issue/3456

---

diff --git a/cli/koji_cli/commands.py b/cli/koji_cli/commands.py
index 0a4e7e4..cf76ccc 100644
--- a/cli/koji_cli/commands.py
+++ b/cli/koji_cli/commands.py
@@ -6952,50 +6952,49 @@ def anon_handle_download_task(options, session, args):
     for task in list_tasks:
         taskarch = task['arch']
         task_id = str(task['id'])
-        if len(suboptions.arches) == 0 or taskarch in suboptions.arches:
-            files = list_task_output_all_volumes(session, task["id"])
-            for filename in files:
-                if filename.endswith('src.rpm'):
-                    filetype = 'src.rpm'
-                else:
-                    filetype = filename.rsplit('.', 1)[1]
-                if suboptions.all and not (task['method'] in build_methods_list and
-                                           filetype in rpm_file_types):
-                    if filetype != 'log':
-                        for volume in files[filename]:
-                            if suboptions.dirpertask:
-                                new_filename = '%s/%s' % (task_id, filename)
-                            else:
-                                if taskarch not in filename and filetype != 'src.rpm':
-                                    part_filename = filename[:-len('.%s' % filetype)]
-                                    new_filename = "%s.%s.%s" % (part_filename,
-                                                                 taskarch, filetype)
-                                else:
-                                    new_filename = filename
-                            downloads.append((task, filename, volume, new_filename, task_id))
-                elif task['method'] in build_methods_list:
-                    if filetype in rpm_file_types:
-                        filearch = filename.split(".")[-2]
-                        for volume in files[filename]:
-                            if len(suboptions.arches) == 0 or filearch in suboptions.arches:
-                                if suboptions.dirpertask:
-                                    new_filename = '%s/%s' % (task_id, filename)
-                                else:
-                                    new_filename = filename
-                                downloads.append((task, filename, volume, new_filename,
-                                                  task_id))
-
-                if filetype == 'log' and suboptions.logs:
+        files = list_task_output_all_volumes(session, task["id"])
+        for filename in files:
+            if filename.endswith('src.rpm'):
+                filetype = 'src.rpm'
+            else:
+                filetype = filename.rsplit('.', 1)[1]
+            if suboptions.all and not (task['method'] in build_methods_list and
+                                       filetype in rpm_file_types):
+                if filetype != 'log':
                     for volume in files[filename]:
                         if suboptions.dirpertask:
                             new_filename = '%s/%s' % (task_id, filename)
                         else:
-                            if taskarch not in filename:
-                                part_filename = filename[:-len('.log')]
-                                new_filename = "%s.%s.log" % (part_filename, taskarch)
+                            if taskarch not in filename and filetype != 'src.rpm':
+                                part_filename = filename[:-len('.%s' % filetype)]
+                                new_filename = "%s.%s.%s" % (part_filename,
+                                                             taskarch, filetype)
                             else:
                                 new_filename = filename
                         downloads.append((task, filename, volume, new_filename, task_id))
+            elif task['method'] in build_methods_list:
+                if filetype in rpm_file_types:
+                    filearch = filename.split(".")[-2]
+                    for volume in files[filename]:
+                        if len(suboptions.arches) == 0 or filearch in suboptions.arches:
+                            if suboptions.dirpertask:
+                                new_filename = '%s/%s' % (task_id, filename)
+                            else:
+                                new_filename = filename
+                            downloads.append((task, filename, volume, new_filename,
+                                              task_id))
+
+            if filetype == 'log' and suboptions.logs:
+                for volume in files[filename]:
+                    if suboptions.dirpertask:
+                        new_filename = '%s/%s' % (task_id, filename)
+                    else:
+                        if taskarch not in filename:
+                            part_filename = filename[:-len('.log')]
+                            new_filename = "%s.%s.log" % (part_filename, taskarch)
+                        else:
+                            new_filename = filename
+                    downloads.append((task, filename, volume, new_filename, task_id))
 
     if len(downloads) == 0:
         print("No files for download found.")
diff --git a/tests/test_cli/test_download_task.py b/tests/test_cli/test_download_task.py
index a399810..33b1667 100644
--- a/tests/test_cli/test_download_task.py
+++ b/tests/test_cli/test_download_task.py
@@ -120,26 +120,29 @@ Default behavior without --all option downloads .rpm files only for build and bu
     def test_handle_download_task_parent(self):
         args = [str(self.parent_task_id), '--arch=noarch,x86_64']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2},
-                                                     {'id': 33333,
-                                                      'method': 'buildArch',
-                                                      'arch': 'x86_64',
-                                                      'state': 2},
-                                                     {'id': 44444,
-                                                      'method': 'buildArch',
-                                                      'arch': 's390',
-                                                      'state': 2},
-                                                     {'id': 55555,
-                                                      'method': 'tagBuild',
-                                                      'arch': 'noarch',
-                                                      'state': 2}
-                                                     ]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2},
+            {'id': 33333,
+             'method': 'buildArch',
+             'arch': 'x86_64',
+             'state': 2},
+            {'id': 44444,
+             'method': 'buildArch',
+             'arch': 's390',
+             'state': 2},
+            {'id': 55555,
+             'method': 'tagBuild',
+             'arch': 'noarch',
+             'state': 2}
+        ]
         self.list_task_output_all_volumes.side_effect = [
+            {},
             {'somerpm.src.rpm': ['DEFAULT', 'vol1']},
             {'somerpm.x86_64.rpm': ['DEFAULT', 'vol2']},
+            {'somerpm.s390.rpm': ['DEFAULT']},
             {'somerpm.noarch.rpm': ['vol3'],
              'somelog.log': ['DEFAULT', 'vol1']},
         ]
@@ -156,8 +159,10 @@ Default behavior without --all option downloads .rpm files only for build and bu
         self.session.getTaskInfo.assert_called_once_with(self.parent_task_id)
         self.session.getTaskChildren.assert_called_once_with(self.parent_task_id)
         self.assertEqual(self.list_task_output_all_volumes.mock_calls, [
+            call(self.session, 123333),
             call(self.session, 22222),
             call(self.session, 33333),
+            call(self.session, 44444),
             call(self.session, 55555)])
         self.assertListEqual(self.download_file.mock_calls, [
             call('https://topurl/work/tasks/3333/33333/somerpm.x86_64.rpm',
@@ -169,10 +174,11 @@ Default behavior without --all option downloads .rpm files only for build and bu
     def test_handle_download_task_log(self):
         args = [str(self.parent_task_id), '--log']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2}]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2}]
         self.list_task_output_all_volumes.side_effect = [{}, {
             'somerpm.src.rpm': ['DEFAULT', 'vol1'],
             'somerpm.x86_64.rpm': ['DEFAULT', 'vol2'],
@@ -229,15 +235,17 @@ Default behavior without --all option downloads .rpm files only for build and bu
         self.ensure_connection.assert_called_once_with(self.session, self.options)
         self.session.getTaskInfo.assert_called_once_with(self.parent_task_id)
         self.session.getTaskChildren.assert_called_once_with(self.parent_task_id)
-        self.list_task_output_all_volumes.assert_not_called()
+        self.list_task_output_all_volumes.assert_called_once_with(
+            self.session, self.parent_task_id)
         self.download_file.assert_not_called()
 
     def test_handle_download_parent_not_finished(self):
         args = [str(self.parent_task_id)]
-        self.session.getTaskInfo.return_value = {'id': self.parent_task_id,
-                                                 'method': 'buildArch',
-                                                 'arch': 'taskarch',
-                                                 'state': 3}
+        self.session.getTaskInfo.return_value = {
+            'id': self.parent_task_id,
+            'method': 'buildArch',
+            'arch': 'taskarch',
+            'state': 3}
         self.list_task_output_all_volumes.return_value = {
             'somerpm.src.rpm': ['DEFAULT', 'vol1'],
             'somerpm.x86_64.rpm': ['DEFAULT', 'vol2'],
@@ -265,10 +273,11 @@ Default behavior without --all option downloads .rpm files only for build and bu
     def test_handle_download_child_not_finished(self):
         args = [str(self.parent_task_id)]
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 3}]
+        self.session.getTaskChildren.return_value = [{
+            'id': 22222,
+            'method': 'buildArch',
+            'arch': 'noarch',
+            'state': 3}]
         self.list_task_output_all_volumes.side_effect = [
             {'somerpm.src.rpm': ['DEFAULT', 'vol1']},
             {'somenextrpm.src.rpm': ['DEFAULT', 'vol1']}]
@@ -374,26 +383,29 @@ Options:
     def test_handle_download_task_parent_dirpertask(self):
         args = [str(self.parent_task_id), '--arch=noarch,x86_64', '--dirpertask', '--log']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2},
-                                                     {'id': 33333,
-                                                      'method': 'buildArch',
-                                                      'arch': 'x86_64',
-                                                      'state': 2},
-                                                     {'id': 44444,
-                                                      'method': 'buildArch',
-                                                      'arch': 's390',
-                                                      'state': 2},
-                                                     {'id': 55555,
-                                                      'method': 'tagBuild',
-                                                      'arch': 'noarch',
-                                                      'state': 2}
-                                                     ]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2},
+            {'id': 33333,
+             'method': 'buildArch',
+             'arch': 'x86_64',
+             'state': 2},
+            {'id': 44444,
+             'method': 'buildArch',
+             'arch': 's390',
+             'state': 2},
+            {'id': 55555,
+             'method': 'tagBuild',
+             'arch': 'noarch',
+             'state': 2}
+        ]
         self.list_task_output_all_volumes.side_effect = [
+            {},
             {'somerpm.src.rpm': ['DEFAULT', 'vol1'], 'somerpm.noarch.rpm': ['DEFAULT']},
             {'somerpm.x86_64.rpm': ['DEFAULT', 'vol2']},
+            {'somerpm.s390.rpm': ['DEFAULT']},
             {'somerpm.noarch.rpm': ['vol3'],
              'somelog.log': ['DEFAULT', 'vol1']},
         ]
@@ -410,8 +422,10 @@ Options:
         self.session.getTaskInfo.assert_called_once_with(self.parent_task_id)
         self.session.getTaskChildren.assert_called_once_with(self.parent_task_id)
         self.assertEqual(self.list_task_output_all_volumes.mock_calls, [
+            call(self.session, 123333),
             call(self.session, 22222),
             call(self.session, 33333),
+            call(self.session, 44444),
             call(self.session, 55555)])
         self.assertListEqual(self.download_file.mock_calls, [
             call('https://topurl/work/tasks/2222/22222/somerpm.noarch.rpm',
@@ -430,26 +444,29 @@ Options:
     def test_handle_download_task_parent_all(self):
         args = [str(self.parent_task_id), '--arch=noarch,x86_64', '--all']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2},
-                                                     {'id': 33333,
-                                                      'method': 'buildArch',
-                                                      'arch': 'x86_64',
-                                                      'state': 2},
-                                                     {'id': 44444,
-                                                      'method': 'buildArch',
-                                                      'arch': 's390',
-                                                      'state': 2},
-                                                     {'id': 55555,
-                                                      'method': 'tagBuild',
-                                                      'arch': 'noarch',
-                                                      'state': 2}
-                                                     ]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2},
+            {'id': 33333,
+             'method': 'buildArch',
+             'arch': 'x86_64',
+             'state': 2},
+            {'id': 44444,
+             'method': 'buildArch',
+             'arch': 's390',
+             'state': 2},
+            {'id': 55555,
+             'method': 'tagBuild',
+             'arch': 'noarch',
+             'state': 2}
+        ]
         self.list_task_output_all_volumes.side_effect = [
+            {},
             {'somerpm.src.rpm': ['DEFAULT', 'vol1']},
             {'somerpm.json': ['DEFAULT', 'vol2']},
+            {'somerpm.s390.rpm': ['DEFAULT']},
             {'somerpm.noarch.rpm': ['vol3'],
              'somelog.log': ['DEFAULT', 'vol1']},
         ]
@@ -466,8 +483,10 @@ Options:
         self.session.getTaskInfo.assert_called_once_with(self.parent_task_id)
         self.session.getTaskChildren.assert_called_once_with(self.parent_task_id)
         self.assertEqual(self.list_task_output_all_volumes.mock_calls, [
+            call(self.session, 123333),
             call(self.session, 22222),
             call(self.session, 33333),
+            call(self.session, 44444),
             call(self.session, 55555)])
         self.assertListEqual(self.download_file.mock_calls, [
             call('https://topurl/work/tasks/3333/33333/somerpm.json',
@@ -504,23 +523,24 @@ Options:
     def test_handle_download_task_parent_dirpertask_all(self):
         args = [str(self.parent_task_id), '--dirpertask', '--all']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2},
-                                                     {'id': 33333,
-                                                      'method': 'buildArch',
-                                                      'arch': 'x86_64',
-                                                      'state': 2},
-                                                     {'id': 44444,
-                                                      'method': 'buildArch',
-                                                      'arch': 's390',
-                                                      'state': 2},
-                                                     {'id': 55555,
-                                                      'method': 'tagBuild',
-                                                      'arch': 'noarch',
-                                                      'state': 2}
-                                                     ]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2},
+            {'id': 33333,
+             'method': 'buildArch',
+             'arch': 'x86_64',
+             'state': 2},
+            {'id': 44444,
+             'method': 'buildArch',
+             'arch': 's390',
+             'state': 2},
+            {'id': 55555,
+             'method': 'tagBuild',
+             'arch': 'noarch',
+             'state': 2}
+        ]
         self.list_task_output_all_volumes.side_effect = [
             {},
             {'somerpm.src.rpm': ['DEFAULT', 'vol1'], 'somerpm.noarch.rpm': ['DEFAULT']},
@@ -568,11 +588,12 @@ Options:
     def test_handle_download_task_log_with_arch(self):
         args = [str(self.parent_task_id), '--arch=noarch', '--log']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2}]
-        self.list_task_output_all_volumes.side_effect = [{
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2}]
+        self.list_task_output_all_volumes.side_effect = [{}, {
             'somerpm.src.rpm': ['DEFAULT', 'vol1'],
             'somerpm.x86_64.rpm': ['DEFAULT', 'vol2'],
             'somerpm.noarch.rpm': ['vol3'],
@@ -604,23 +625,24 @@ Options:
     def test_handle_download_task_all_log(self):
         args = [str(self.parent_task_id), '--all', '--log']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2},
-                                                     {'id': 33333,
-                                                      'method': 'buildArch',
-                                                      'arch': 'x86_64',
-                                                      'state': 2},
-                                                     {'id': 44444,
-                                                      'method': 'buildArch',
-                                                      'arch': 's390',
-                                                      'state': 2},
-                                                     {'id': 55555,
-                                                      'method': 'tagBuild',
-                                                      'arch': 'noarch',
-                                                      'state': 2}
-                                                     ]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2},
+            {'id': 33333,
+             'method': 'buildArch',
+             'arch': 'x86_64',
+             'state': 2},
+            {'id': 44444,
+             'method': 'buildArch',
+             'arch': 's390',
+             'state': 2},
+            {'id': 55555,
+             'method': 'tagBuild',
+             'arch': 'noarch',
+             'state': 2}
+        ]
         self.list_task_output_all_volumes.side_effect = [
             {},
             {'somerpm.src.rpm': ['DEFAULT', 'vol1'], 'somerpm.noarch.rpm': ['DEFAULT']},
@@ -671,23 +693,24 @@ Options:
     def test_handle_download_task_parent_dirpertask_all_conflict_names(self):
         args = [str(self.parent_task_id), '--all']
         self.session.getTaskInfo.return_value = self.parent_task_info
-        self.session.getTaskChildren.return_value = [{'id': 22222,
-                                                      'method': 'buildArch',
-                                                      'arch': 'noarch',
-                                                      'state': 2},
-                                                     {'id': 33333,
-                                                      'method': 'buildArch',
-                                                      'arch': 'x86_64',
-                                                      'state': 2},
-                                                     {'id': 44444,
-                                                      'method': 'buildArch',
-                                                      'arch': 's390',
-                                                      'state': 2},
-                                                     {'id': 55555,
-                                                      'method': 'tagBuild',
-                                                      'arch': 'noarch',
-                                                      'state': 2}
-                                                     ]
+        self.session.getTaskChildren.return_value = [
+            {'id': 22222,
+             'method': 'buildArch',
+             'arch': 'noarch',
+             'state': 2},
+            {'id': 33333,
+             'method': 'buildArch',
+             'arch': 'x86_64',
+             'state': 2},
+            {'id': 44444,
+             'method': 'buildArch',
+             'arch': 's390',
+             'state': 2},
+            {'id': 55555,
+             'method': 'tagBuild',
+             'arch': 'noarch',
+             'state': 2}
+        ]
         self.list_task_output_all_volumes.side_effect = [
             {},
             {'somerpm.src.rpm': ['DEFAULT', 'vol1'], 'somerpm.noarch.rpm': ['DEFAULT']},

