Merge pull request #5 from lbryio/fix-unclean-exit

Fix unclean exit due to uncaught errors/exceptions

aioupnp/gateway.py

@@ -155,8 +155,8 @@ class Gateway:
         }
 
     @classmethod
-    async def _discover_gateway(cls, lan_address: str, gateway_address: str, timeout: int=30,
+    async def _discover_gateway(cls, lan_address: str, gateway_address: str, timeout: int = 30,
-                                igd_args: OrderedDict=None, loop=None, unicast: bool=False):
+                                igd_args: OrderedDict = None, loop=None, unicast: bool = False):
         ignored: set = set()
         required_commands = [
             'AddPortMapping',
@@ -196,19 +196,26 @@ class Gateway:
     async def discover_gateway(cls, lan_address: str, gateway_address: str, timeout: int = 30,
                                igd_args: OrderedDict = None, loop=None, unicast: bool = None):
         if unicast is not None:
-            return await cls._discover_gateway(lan_address, gateway_address, timeout, igd_args, loop)
+            return await cls._discover_gateway(lan_address, gateway_address, timeout, igd_args, loop, unicast)
+
         done, pending = await asyncio.wait([
             cls._discover_gateway(
                 lan_address, gateway_address, timeout, igd_args, loop, unicast=True
             ),
             cls._discover_gateway(
                 lan_address, gateway_address, timeout, igd_args, loop, unicast=False
-            )], return_when=asyncio.tasks.FIRST_COMPLETED
             )
-        for task in list(pending):
+        ], return_when=asyncio.tasks.FIRST_COMPLETED)
+
+        for task in pending:
             task.cancel()
-        result = list(done)[0].result()
+        for task in done:
-        return result
+            try:
+                task.exception()
+            except asyncio.CancelledError:
+                pass
+
+        return list(done)[0].result()
 
     async def discover_commands(self, loop=None):
         response, xml_bytes, get_err = await scpd_get(self.path.decode(), self.base_ip.decode(), self.port, loop=loop)

aioupnp/upnp.py

@@ -394,7 +394,7 @@ class UPnP:
         try:
             result = fut.result()
         except UPnPError as err:
-            print("aioupnp encountered an error:\n%s" % str(err))
+            print("aioupnp encountered an error: %s" % str(err))
             return
 
         if isinstance(result, (list, tuple, dict)):

tests/test_cli.py

@@ -85,7 +85,7 @@ class TestCLI(TestBase):
 
     def test_m_search(self):
         actual_output = StringIO()
-        timeout_msg = "aioupnp encountered an error:\nM-SEARCH for 10.0.0.1:1900 timed out\n"
+        timeout_msg = "aioupnp encountered an error: M-SEARCH for 10.0.0.1:1900 timed out\n"
         with contextlib.redirect_stdout(actual_output):
             with mock_tcp_and_udp(self.loop, '10.0.0.1', tcp_replies=self.scpd_replies, udp_replies=self.udp_replies):
                 main(

tests/test_gateway.py

@@ -1,3 +1,6 @@
+import asyncio
+
+from aioupnp.fault import UPnPError
 from tests import TestBase
 from tests.mocks import mock_tcp_and_udp
 from collections import OrderedDict
@@ -62,6 +65,16 @@ class TestDiscoverDLinkDIR890L(TestBase):
         'GetExternalIPAddress': 'urn:schemas-upnp-org:service:WANIPConnection:1'
     }
 
+    async def test_discover_gateway(self):
+        with self.assertRaises(UPnPError) as e1:
+            with mock_tcp_and_udp(self.loop):
+                await Gateway.discover_gateway("10.0.0.2", "10.0.0.1", 2)
+        with self.assertRaises(UPnPError) as e2:
+            with mock_tcp_and_udp(self.loop):
+                await Gateway.discover_gateway("10.0.0.2", "10.0.0.1", 2, unicast=False)
+        self.assertEqual(str(e1.exception), "M-SEARCH for 10.0.0.1:1900 timed out")
+        self.assertEqual(str(e2.exception), "M-SEARCH for 10.0.0.1:1900 timed out")
+
     async def test_discover_commands(self):
         with mock_tcp_and_udp(self.loop, tcp_replies=self.replies):
             gateway = Gateway(self.reply, self.m_search_args, self.client_address, self.gateway_address)