From 9b8caf9e177e6bc1d4a4ac0350412e61733d58f7 Mon Sep 17 00:00:00 2001
From: Viktor Stanchev <me@viktorstanchev.com>
Date: Sat, 19 Sep 2015 22:27:05 -0700
Subject: [PATCH] fix mirror sorting, test mirror choice

---
 pirate/pirate.py     | 28 ++++++++++++++++------------
 tests/test_pirate.py | 22 ++++++++++++++++++++++
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/pirate/pirate.py b/pirate/pirate.py
index 8c5963b..e657573 100755
--- a/pirate/pirate.py
+++ b/pirate/pirate.py
@@ -7,6 +7,7 @@ import socket
 import urllib.request as request
 import urllib.error
 import sys
+from collections import OrderedDict
 
 import webbrowser
 
@@ -226,8 +227,9 @@ def combine_configs(config, args):
     return args
 
 
-def search_mirrors(args):
-    mirrors = {'https://thepiratebay.mn'}
+def search_mirrors(pages, category, sort, action, search):
+    mirrors = OrderedDict()
+    mirrors['https://thepiratebay.mn'] = None
     try:
         req = request.Request('https://proxybay.co/list.txt',
                               headers=pirate.data.default_headers)
@@ -237,19 +239,21 @@ def search_mirrors(args):
     else:
         if f.getcode() != 200:
             raise IOError('The proxy bay responded with an error.')
-        mirrors = mirrors.union([i.decode('utf-8').strip()
-                                for i in f.readlines()][3:]
-                                ).difference(pirate.data.blacklist)
+        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:
+    for mirror in mirrors.keys():
         try:
             print('Trying', mirror, end='... \n')
             results = pirate.torrent.remote(
-                pages=args.pages,
-                category=pirate.torrent.parse_category(args.category),
-                sort=pirate.torrent.parse_sort(args.sort),
-                mode=args.action,
-                terms=args.search,
+                pages=pages,
+                category=pirate.torrent.parse_category(category),
+                sort=pirate.torrent.parse_sort(sort),
+                mode=action,
+                terms=search,
                 mirror=mirror
             )
         except (urllib.error.URLError, socket.timeout,
@@ -296,7 +300,7 @@ def main():
     if args.source == 'local_tpb':
         results = pirate.local.search(args.database, args.search)
     elif args.source == 'tpb':
-        results, site = search_mirrors(args)
+        results, site = search_mirrors(args.pages, args.category, args.sort, args.action, args.search)
 
     if len(results) == 0:
         print('No results')
diff --git a/tests/test_pirate.py b/tests/test_pirate.py
index a7f51df..5f6d45c 100755
--- a/tests/test_pirate.py
+++ b/tests/test_pirate.py
@@ -1,5 +1,9 @@
 #!/usr/bin/env python3
+import socket
 import unittest
+from unittest import mock
+from unittest.mock import patch, call
+
 import pirate.pirate
 
 
@@ -108,5 +112,23 @@ class TestPirate(unittest.TestCase):
                 value = getattr(args, option)
                 self.assertEqual(test[2][option], value)
 
+    def test_search_mirrors(self):
+        pages, category, sort, action, search = (1, 100, 10, 'browse', [])
+        class MockResponse():
+            readlines = mock.MagicMock(return_value=[x.encode('utf-8') for x in ['', '', '', 'https://example.com']])
+            info = mock.MagicMock()
+            getcode = mock.MagicMock(return_value=200)
+        response_obj = MockResponse()
+        with patch('urllib.request.urlopen', return_value=response_obj) as urlopen:
+            with patch('pirate.torrent.remote', return_value=[]) as remote:
+                results, mirror = pirate.pirate.search_mirrors(pages, category, sort, action, search)
+                remote.assert_called_once_with(pages=1, category=100, sort=10, mode='browse', terms=[], mirror='https://thepiratebay.mn')
+            with patch('pirate.torrent.remote', side_effect=[socket.timeout, []]) as remote:
+                results, mirror = pirate.pirate.search_mirrors(pages, category, sort, action, search)
+                remote.assert_has_calls([
+                    call(pages=1, category=100, sort=10, mode='browse', terms=[], mirror='https://thepiratebay.mn'),
+                    call(pages=1, category=100, sort=10, mode='browse', terms=[], mirror='https://example.com')
+                ])
+
 if __name__ == '__main__':
     unittest.main()