qvm-template: Add tests for download function and fix minor bugs

qubesadmin/tests/tools/qvm_template.py

@@ -191,7 +191,7 @@ class TC_00_qvm_template(qubesadmin.tests.QubesTestCase):
         mock_dl_list.return_value = {}
         mock_call.side_effect = self.add_new_vm_side_effect
         mock_time = mock.Mock(wraps=datetime.datetime)
-        mock_time.today.return_value = \
+        mock_time.now.return_value = \
             datetime.datetime(2020, 9, 1, 15, 30, tzinfo=datetime.timezone.utc)
         with mock.patch('builtins.open', mock.mock_open()) as mock_open, \
                 mock.patch('datetime.datetime', new=mock_time), \
@@ -322,7 +322,7 @@ class TC_00_qvm_template(qubesadmin.tests.QubesTestCase):
         mock_dl_list.return_value = {}
         mock_call.side_effect = self.add_new_vm_side_effect
         mock_time = mock.Mock(wraps=datetime.datetime)
-        mock_time.today.return_value = \
+        mock_time.now.return_value = \
             datetime.datetime(2020, 9, 1, 15, 30, tzinfo=datetime.timezone.utc)
         with mock.patch('builtins.open', mock.mock_open()) as mock_open, \
                 mock.patch('datetime.datetime', new=mock_time), \
@@ -3243,3 +3243,288 @@ test-vm : Qubes template for fedora-31
             mock.call(args, self.app)
         ])
         self.assertAllCalled()
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_180_download_success(self, mock_qrexec, mock_dllist):
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=1
+            )
+            qubesadmin.tools.qvm_template.download(args, self.app, dir, {
+                    'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                        ('1', '2', '3'), 'qubes-templates-itl', 1048576),
+                    'fedora-32': qubesadmin.tools.qvm_template.DlEntry(
+                        ('0', '1', '2'),
+                        'qubes-templates-itl-testing',
+                        2048576)
+                }, '.unverified')
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm.unverified',
+                    1048576),
+                mock.call(args, self.app, 'qubes-template-fedora-32-0:1-2',
+                    dir + '/qubes-template-fedora-32-0:1-2.rpm.unverified',
+                    2048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])
+            self.assertTrue(all(
+                [x.endswith('.unverified') for x in os.listdir(dir)]))
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_181_download_success_nosuffix(self, mock_qrexec, mock_dllist):
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=1,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err:
+                qubesadmin.tools.qvm_template.download(args, self.app, None, {
+                        'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                            ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                    })
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_182_download_success_getdllist(self, mock_qrexec, mock_dllist):
+        mock_dllist.return_value = {
+            'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+        }
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=1
+            )
+            qubesadmin.tools.qvm_template.download(args, self.app,
+                dir, None, '.unverified',
+                qubesadmin.tools.qvm_template.VersionSelector.LATEST_LOWER)
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm.unverified',
+                    1048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [
+                mock.call(args, self.app,
+                    version_selector=\
+                        qubesadmin.tools.qvm_template.\
+                        VersionSelector.LATEST_LOWER)
+            ])
+            self.assertTrue(all(
+                [x.endswith('.unverified') for x in os.listdir(dir)]))
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_183_download_success_downloaddir(self, mock_qrexec, mock_dllist):
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=1,
+                downloaddir=dir
+            )
+            qubesadmin.tools.qvm_template.download(args, self.app, None, {
+                    'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                        ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                }, '.unverified')
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm.unverified',
+                    1048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])
+            self.assertTrue(all(
+                [x.endswith('.unverified') for x in os.listdir(dir)]))
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_184_download_success_exists(self, mock_qrexec, mock_dllist):
+        with tempfile.TemporaryDirectory() as dir:
+            with open(os.path.join(
+                        dir, 'qubes-template-fedora-31-1:2-3.rpm.unverified'),
+                    'w') as _:
+                pass
+            args = argparse.Namespace(
+                retries=1,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err:
+                qubesadmin.tools.qvm_template.download(args, self.app, None, {
+                        'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                            ('1', '2', '3'), 'qubes-templates-itl', 1048576),
+                        'fedora-32': qubesadmin.tools.qvm_template.DlEntry(
+                            ('0', '1', '2'),
+                            'qubes-templates-itl-testing',
+                            2048576)
+                    }, '.unverified')
+                self.assertTrue('already exists, skipping'
+                    in mock_err.getvalue())
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-32-0:1-2',
+                    dir + '/qubes-template-fedora-32-0:1-2.rpm.unverified',
+                    2048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])
+            self.assertTrue(all(
+                [x.endswith('.unverified') for x in os.listdir(dir)]))
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_185_download_success_existsmove(self, mock_qrexec, mock_dllist):
+        with tempfile.TemporaryDirectory() as dir:
+            with open(os.path.join(
+                        dir, 'qubes-template-fedora-31-1:2-3.rpm'),
+                    'w') as _:
+                pass
+            args = argparse.Namespace(
+                retries=1,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err:
+                qubesadmin.tools.qvm_template.download(args, self.app, None, {
+                        'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                            ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                    }, '.unverified')
+                self.assertTrue('already exists, skipping'
+                    in mock_err.getvalue())
+            self.assertEqual(mock_qrexec.mock_calls, [])
+            self.assertEqual(mock_dllist.mock_calls, [])
+            self.assertTrue(os.path.exists(
+                dir + '/qubes-template-fedora-31-1:2-3.rpm.unverified'))
+            self.assertTrue(all(
+                [x.endswith('.unverified') for x in os.listdir(dir)]))
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_186_download_success_existsnosuffix(self, mock_qrexec, mock_dllist):
+        with tempfile.TemporaryDirectory() as dir:
+            with open(os.path.join(
+                        dir, 'qubes-template-fedora-31-1:2-3.rpm'),
+                    'w') as _:
+                pass
+            args = argparse.Namespace(
+                retries=1,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err:
+                qubesadmin.tools.qvm_template.download(args, self.app, None, {
+                        'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                            ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                    })
+                self.assertTrue('already exists, skipping'
+                    in mock_err.getvalue())
+            self.assertEqual(mock_qrexec.mock_calls, [])
+            self.assertEqual(mock_dllist.mock_calls, [])
+            self.assertTrue(os.path.exists(
+                dir + '/qubes-template-fedora-31-1:2-3.rpm'))
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_187_download_success_retry(self, mock_qrexec, mock_dllist):
+        counter = 0
+        def f(*args):
+            nonlocal counter
+            counter += 1
+            if counter == 1:
+                raise ConnectionError
+        mock_qrexec.side_effect = f
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=2,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err, \
+                    mock.patch('os.remove') as mock_rm:
+                qubesadmin.tools.qvm_template.download(args, self.app, None, {
+                        'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                            ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                    })
+                self.assertTrue('retrying...' in mock_err.getvalue())
+                self.assertEqual(mock_rm.mock_calls, [
+                    mock.call(dir + '/qubes-template-fedora-31-1:2-3.rpm')
+                ])
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576),
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_188_download_fail_retry(self, mock_qrexec, mock_dllist):
+        counter = 0
+        def f(*args):
+            nonlocal counter
+            counter += 1
+            if counter <= 3:
+                raise ConnectionError
+        mock_qrexec.side_effect = f
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=3,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err, \
+                    mock.patch('os.remove') as mock_rm:
+                with self.assertRaises(SystemExit):
+                    qubesadmin.tools.qvm_template.download(
+                        args, self.app, None, {
+                            'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                                ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                        })
+                self.assertEqual(mock_err.getvalue().count('retrying...'), 2)
+                self.assertTrue('download failed' in mock_err.getvalue())
+                self.assertEqual(mock_rm.mock_calls, [
+                    mock.call(dir + '/qubes-template-fedora-31-1:2-3.rpm'),
+                    mock.call(dir + '/qubes-template-fedora-31-1:2-3.rpm'),
+                    mock.call(dir + '/qubes-template-fedora-31-1:2-3.rpm')
+                ])
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576),
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576),
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])
+
+    @mock.patch('qubesadmin.tools.qvm_template.get_dl_list')
+    @mock.patch('qubesadmin.tools.qvm_template.qrexec_download')
+    def test_189_download_fail_interrupt(self, mock_qrexec, mock_dllist):
+        def f(*args):
+            raise RuntimeError
+        mock_qrexec.side_effect = f
+        with tempfile.TemporaryDirectory() as dir:
+            args = argparse.Namespace(
+                retries=3,
+                downloaddir=dir
+            )
+            with mock.patch('sys.stderr', new=io.StringIO()) as mock_err, \
+                    mock.patch('os.remove') as mock_rm:
+                with self.assertRaises(RuntimeError):
+                    qubesadmin.tools.qvm_template.download(
+                        args, self.app, None, {
+                            'fedora-31': qubesadmin.tools.qvm_template.DlEntry(
+                                ('1', '2', '3'), 'qubes-templates-itl', 1048576)
+                        })
+                self.assertEqual(mock_rm.mock_calls, [
+                    mock.call(dir + '/qubes-template-fedora-31-1:2-3.rpm')
+                ])
+            self.assertEqual(mock_qrexec.mock_calls, [
+                mock.call(args, self.app, 'qubes-template-fedora-31-1:2-3',
+                    dir + '/qubes-template-fedora-31-1:2-3.rpm',
+                    1048576)
+            ])
+            self.assertEqual(mock_dllist.mock_calls, [])

qubesadmin/tools/qvm_template.py

@@ -53,6 +53,8 @@ def qubes_release() -> str:
                 continue
             val = val.strip('\'"') # strip possible quotes
             return val
+    # Return default value instead of throwing so that it works on CI
+    return '4.1'
 
 def parser_gen() -> argparse.ArgumentParser:
     """Generate argument parser for the application."""
@@ -257,7 +259,7 @@ def is_match_spec(name: str, epoch: str, version: str, release: str, spec: str
     :return: A tuple. The first element indicates whether there is a match; the
         second element represents the priority of the match (lower is better)
     """
-    if epoch != 0:
+    if epoch != '0':
         targets = [
             f'{name}-{epoch}:{version}-{release}',
             f'{name}',
@@ -713,13 +715,12 @@ def download(
         spec = PACKAGE_NAME_PREFIX + name + '-' + version_str
         target = os.path.join(path, '%s.rpm' % spec)
         target_suffix = target + suffix
-        if suffix != '' and os.path.exists(target_suffix):
+        if os.path.exists(target_suffix):
             print('\'%s\' already exists, skipping...' % target,
                 file=sys.stderr)
-        if os.path.exists(target):
+        elif os.path.exists(target):
             print('\'%s\' already exists, skipping...' % target,
                 file=sys.stderr)
-            if suffix != '':
             os.rename(target, target_suffix)
         else:
             print('Downloading \'%s\'...' % spec, file=sys.stderr)
@@ -930,7 +931,7 @@ def install(
                             tz=datetime.timezone.utc) \
                         .strftime(DATE_FMT)
                 tpl.features['template-installtime'] = \
-                    datetime.datetime.today(
+                    datetime.datetime.now(
                         tz=datetime.timezone.utc).strftime(DATE_FMT)
                 tpl.features['template-license'] = \
                     package_hdr[rpm.RPMTAG_LICENSE]
@@ -1100,18 +1101,19 @@ def list_templates(args: argparse.Namespace,
 
     if args.machine_readable:
         if operation == 'info':
-            tpl_list = info_to_machine_output(tpl_list)
+            tpl_list_dict = info_to_machine_output(tpl_list)
         elif operation == 'list':
-            tpl_list = list_to_machine_output(tpl_list)
+            tpl_list_dict = list_to_machine_output(tpl_list)
-        for status, grp in tpl_list.items():
+        for status, grp in tpl_list_dict.items():
             for line in grp:
                 print('|'.join([status] + list(line.values())))
     elif args.machine_readable_json:
         if operation == 'info':
-            tpl_list = info_to_machine_output(tpl_list, replace_newline=False)
+            tpl_list_dict = \
+                info_to_machine_output(tpl_list, replace_newline=False)
         elif operation == 'list':
-            tpl_list = list_to_machine_output(tpl_list)
+            tpl_list_dict = list_to_machine_output(tpl_list)
-        print(json.dumps(tpl_list))
+        print(json.dumps(tpl_list_dict))
     else:
         if operation == 'info':
             tpl_list = info_to_human_output(tpl_list)