Merge pull request #88 from vikstrous/search-mirror

tidy up search_mirrors function

pirate/data.py

@@ -11,3 +11,6 @@ blacklist = set(json.loads(get_resource('blacklist.json').decode()))
 
 default_headers = {'User-Agent': 'pirate get'}
 default_timeout = 10
+
+default_mirror = 'https://thepiratebay.org/'
+mirror_list = 'https://proxybay.co/list.txt'

pirate/pirate.py

@@ -228,33 +228,9 @@ def combine_configs(config, args):
     return args
 
 
-def search_mirrors(printer, pages, category, sort, action, search):
-    mirror_sources = [None, 'https://proxybay.co/list.txt']
-    for mirror_source in mirror_sources:
-        mirrors = OrderedDict()
-        if mirror_source is None:
-            mirrors['https://thepiratebay.mn'] = None
-        else:
+def connect_mirror(mirror, printer, pages, category, sort, action, search):
     try:
-                req = request.Request(mirror_source,
-                                      headers=pirate.data.default_headers)
-                f = request.urlopen(req, timeout=pirate.data.default_timeout)
-            except IOError:
-                printer.print('Could not fetch additional mirrors',
-                              color='WARN')
-            else:
-                if f.getcode() != 200:
-                    raise IOError('The proxy bay responded with an error.')
-                for mirror in [i.decode('utf-8').strip()
-                               for i in f.readlines()][3:]:
-                    mirrors[mirror] = None
-            for mirror in pirate.data.blacklist:
-                if mirror in mirrors:
-                    del mirrors[mirror]
-
-        for mirror in mirrors.keys():
-            try:
-                printer.print('Trying', mirror, end='... \n')
+        printer.print('Trying', mirror, end='... ')
         results = pirate.torrent.remote(
             printer=printer,
             pages=pages,
@@ -262,17 +238,44 @@ def search_mirrors(printer, pages, category, sort, action, search):
             sort=pirate.torrent.parse_sort(printer, sort),
             mode=action,
             terms=search,
-                    mirror=mirror
-                )
-            except (urllib.error.URLError, socket.timeout,
-                    IOError, ValueError):
+            mirror=mirror)
+    except (urllib.error.URLError, socket.timeout, IOError, ValueError):
         printer.print('Failed', color='WARN')
+        return None
     else:
         printer.print('Ok', color='alt')
         return results, mirror
+
+
+def search_mirrors(printer, *args):
+    # try official site
+    result = connect_mirror(pirate.data.default_mirror, printer, *args)
+    if result:
+        return result
+
+    # download mirror list
+    try:
+        req = request.Request(pirate.data.mirror_list,
+                              headers=pirate.data.default_headers)
+        f = request.urlopen(req, timeout=pirate.data.default_timeout)
+    except urllib.error.URLError as e:
+        raise IOError('Could not fetch mirrors', e.reason)
+
+    if f.getcode() != 200:
+        raise IOError('The proxy bay responded with an error',
+                      f.read().decode('utf-8'))
+
+    mirrors = [i.decode('utf-8').strip() for i in f.readlines()][3:]
+
+    # try mirrors
+    for mirror in mirrors:
+        if mirror in pirate.data.blacklist:
+            continue
+        result = connect_mirror(mirror, printer, *args)
+        if result:
+            return result
     else:
-        printer.print('No available mirrors :(', color='WARN')
-        return [], None
+        raise IOError('No more available mirrors')
 
 
 def pirate_main(args):
@@ -308,8 +311,14 @@ def pirate_main(args):
     if args.source == 'local_tpb':
         results = pirate.local.search(args.database, args.search)
     elif args.source == 'tpb':
+        try:
             results, site = search_mirrors(printer, args.pages, args.category,
                                            args.sort, args.action, args.search)
+        except IOError as e:
+            printer.print(e.args[0] + ' :( ', color='ERROR')
+            if len(e.args) > 1:
+                printer.print(e.args[1])
+            sys.exit(1)
 
     if len(results) == 0:
         printer.print('No results')

tests/test_pirate.py

@@ -6,6 +6,7 @@ from unittest import mock
 from unittest.mock import patch, call, MagicMock
 
 import pirate.pirate
+import pirate.data
 from pirate.print import Printer
 
 
@@ -157,14 +158,14 @@ class TestPirate(unittest.TestCase):
             with patch('pirate.torrent.remote', return_value=[]) as remote:
                 results, mirror = pirate.pirate.search_mirrors(printer, pages, category, sort, action, search)
                 self.assertEqual(results, [])
-                self.assertEqual(mirror, 'https://thepiratebay.mn')
-                remote.assert_called_once_with(printer=printer, pages=1, category=100, sort=10, mode='browse', terms=[], mirror='https://thepiratebay.mn')
+                self.assertEqual(mirror, pirate.data.default_mirror)
+                remote.assert_called_once_with(printer=printer, pages=1, category=100, sort=10, mode='browse', terms=[], mirror=pirate.data.default_mirror)
             with patch('pirate.torrent.remote', side_effect=[socket.timeout, []]) as remote:
                 results, mirror = pirate.pirate.search_mirrors(printer, pages, category, sort, action, search)
                 self.assertEqual(results, [])
                 self.assertEqual(mirror, 'https://example.com')
                 remote.assert_has_calls([
-                    call(printer=printer, pages=1, category=100, sort=10, mode='browse', terms=[], mirror='https://thepiratebay.mn'),
+                    call(printer=printer, pages=1, category=100, sort=10, mode='browse', terms=[], mirror=pirate.data.default_mirror),
                     call(printer=printer, pages=1, category=100, sort=10, mode='browse', terms=[], mirror='https://example.com')
                 ])