urlfetcher: move mock code to the testsuite

Signed-off-by: Cole Robinson <crobinso@redhat.com>

setup.py

@@ -417,6 +417,7 @@ class TestBaseCommand(distutils.core.Command):
         self.testfile = None
         self._force_verbose = False
         self._external_coverage = False
+        self._urlfetcher_mock = False
 
     def finalize_options(self):
         if self.only:
@@ -462,6 +463,9 @@ class TestBaseCommand(distutils.core.Command):
         for key, val in self._clistate.items():
             setattr(testsmodule.utils.clistate, key, val)
         testsmodule.setup_logging()
+        if self._urlfetcher_mock:
+            import tests.urlfetcher_mock
+            tests.urlfetcher_mock.setup_mock()
 
         # This makes the test runner report results before exiting from ctrl-c
         unittest.installHandler()
@@ -523,6 +527,7 @@ class TestCommand(TestBaseCommand):
         '''
         Finds all the tests modules in tests/, and runs them.
         '''
+        self._urlfetcher_mock = True
         excludes = ["test_dist.py", "test_urls.py", "test_inject.py"]
         testfiles = self._find_tests_in_dir("tests", excludes)
 

tests/test_cli.py

@@ -853,6 +853,9 @@ c.add_invalid("-c foo --cdrom bar", grep="Cannot specify both -c")  # check for
 c.add_invalid("-c qemu:///system", grep="looks like a libvirt URI")  # error for the ambiguous -c vs --connect
 c.add_invalid("--location /", grep="Error validating install location")  # detect_distro failure
 c.add_invalid("--os-variant foo://bar", grep="Unknown libosinfo ID")  # bad full id
+c.add_invalid("--location http://testsuitefail.com", grep="installable distribution")  # will trigger a particular mock failure
+
+
 
 c = vinst.add_category("single-disk-install", "--nographics --noautoconsole --disk %(EXISTIMG1)s")
 c.add_valid("--hvm --import")  # FV Import install

tests/urlfetcher_mock.py

@@ -0,0 +1,87 @@
+# This work is licensed under the GNU GPLv2 or later.
+# See the COPYING file in the top-level directory.
+#
+
+import ftplib
+import os
+import urllib
+from urllib.request import Request
+
+import requests
+
+from virtinst import log
+
+
+def _in_testsuite():
+    return "VIRTINST_TEST_SUITE" in os.environ
+
+
+def _make_mock_url(url, filesyntax):
+    if url.endswith("treeinfo"):
+        # If the url is requesting treeinfo, give a fake treeinfo from
+        # our testsuite data
+        fn = ("%s/data/cli/fakerhel6tree/.treeinfo" %
+                os.path.abspath(os.path.dirname(__file__)))
+        abspath = os.path.abspath(fn)
+    else:
+        # Otherwise just copy this file
+        abspath = os.path.abspath(__file__)
+
+    if filesyntax:
+        return "file://" + abspath
+    return abspath
+
+
+class _MockRequestsResponse:
+    def __init__(self, url):
+        log.debug("mocking requests session for url=%s", url)
+        fn = _make_mock_url(url, filesyntax=False)
+        self._content = open(fn).read()
+        self.headers = {'content-length': len(self._content)}
+
+    def raise_for_status(self):
+        pass
+    def iter_content(self, *args, **kwargs):
+        dummy = args
+        dummy = kwargs
+        return [self._content.encode("utf-8")]
+
+
+class _MockRequestsSession:
+    def close(self):
+        pass
+    def head(self, url, *args, **kwargs):
+        dummy = args
+        dummy = kwargs
+        return _MockRequestsResponse(url)
+    def get(self, url, *args, **kwargs):
+        dummy = args
+        dummy = kwargs
+        if "testsuitefail" in url:
+            raise RuntimeError("testsuitefail seen, raising mock error")
+        return _MockRequestsResponse(url)
+
+
+class _MockFTPSession:
+    def connect(self, *args, **kwargs):
+        pass
+    def login(self, *args, **kwargs):
+        pass
+    def voidcmd(self, *args, **kwargs):
+        pass
+    def quit(self, *args, **kwargs):
+        pass
+    def size(self, url):
+        path = _make_mock_url(url, filesyntax=False)
+        return os.path.getsize(path)
+
+
+def _MockUrllibRequest(url):
+    url = _make_mock_url(url, filesyntax=True)
+    return Request(url)
+
+
+def setup_mock():
+    requests.Session = _MockRequestsSession
+    ftplib.FTP = _MockFTPSession
+    urllib.request.Request = _MockUrllibRequest

virtinst/install/urlfetcher.py

@@ -18,72 +18,6 @@ import requests
 from ..logger import log
 
 
-##############################
-# Mocking for the test suite #
-##############################
-
-def _in_testsuite():
-    return "VIRTINST_TEST_SUITE" in os.environ
-
-
-def _make_mock_url(url, filesyntax):
-    if url.endswith("treeinfo"):
-        # If the url is requesting treeinfo, give a fake treeinfo from
-        # our testsuite data
-        fn = ("%s/../../tests/data/cli/fakerhel6tree/.treeinfo" %
-                os.path.abspath(os.path.dirname(__file__)))
-        abspath = os.path.abspath(fn)
-    else:
-        # Otherwise just copy this file
-        abspath = os.path.abspath(__file__)
-
-    if filesyntax:
-        return "file://" + abspath
-    return abspath
-
-
-class _MockRequestsResponse:
-    def __init__(self, url):
-        log.debug("mocking requests session for url=%s", url)
-        fn = _make_mock_url(url, filesyntax=False)
-        self._content = open(fn).read()
-        self.headers = {'content-length': len(self._content)}
-
-    def raise_for_status(self):
-        pass
-    def iter_content(self, *args, **kwargs):
-        dummy = args
-        dummy = kwargs
-        return [self._content.encode("utf-8")]
-
-
-class _MockRequestsSession:
-    def close(self):
-        pass
-    def head(self, url, *args, **kwargs):
-        dummy = args
-        dummy = kwargs
-        return _MockRequestsResponse(url)
-    def get(self, url, *args, **kwargs):
-        dummy = args
-        dummy = kwargs
-        return _MockRequestsResponse(url)
-
-
-class _MockFTPSession:
-    def connect(self, *args, **kwargs):
-        pass
-    def login(self, *args, **kwargs):
-        pass
-    def voidcmd(self, *args, **kwargs):
-        pass
-    def quit(self, *args, **kwargs):
-        pass
-    def size(self, url):
-        path = _make_mock_url(url, filesyntax=False)
-        return os.path.getsize(path)
-
-
 ###########################
 # Fetcher implementations #
 ###########################
@@ -224,10 +158,10 @@ class _URLFetcher(object):
             self._grabURL(filename, fileobj, fullurl=fullurl)
             log.debug("Saved file to %s", fn)
             return fn
-        except:  # noqa
-            if fn and os.path.exists(fn):  # pragma: no cover
-                os.unlink(fn)  # pragma: no cover
-            raise  # pragma: no cover
+        except BaseException:  # pragma: no cover
+            if fn and os.path.exists(fn):
+                os.unlink(fn)
+            raise
 
     def acquireFileContent(self, filename):
         """
@@ -242,10 +176,7 @@ class _HTTPURLFetcher(_URLFetcher):
     _session = None
 
     def _prepare(self):
-        if _in_testsuite():
-            self._session = _MockRequestsSession()
-        else:
-            self._session = requests.Session()
+        self._session = requests.Session()
 
     def _cleanup(self):
         if self._session:
@@ -305,10 +236,7 @@ class _FTPURLFetcher(_URLFetcher):
 
         try:
             parsed = urllib.parse.urlparse(self.location)
-            if _in_testsuite():
-                self._ftp = _MockFTPSession()
-            else:
-                self._ftp = ftplib.FTP()
+            self._ftp = ftplib.FTP()
             username = urllib.parse.unquote(parsed.username or '')
             password = urllib.parse.unquote(parsed.password or '')
             self._ftp.connect(parsed.hostname, parsed.port or 0)
@@ -323,8 +251,6 @@ class _FTPURLFetcher(_URLFetcher):
         """
         Use urllib and ftplib to grab the file
         """
-        if _in_testsuite():
-            url = _make_mock_url(url, filesyntax=True)
         request = urllib.request.Request(url)
         urlobj = urllib.request.urlopen(request)
         size = self._ftp.size(urllib.parse.urlparse(url)[2])
@@ -349,7 +275,7 @@ class _FTPURLFetcher(_URLFetcher):
             try:
                 # If it's a file
                 self._ftp.size(path)
-            except ftplib.all_errors:
+            except ftplib.all_errors:  # pragma: no cover
                 # If it's a dir
                 self._ftp.cwd(path)
         except ftplib.all_errors as e:  # pragma: no cover