Use done and pending, add unit test

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',
@@ -197,26 +197,25 @@ class Gateway:
                                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, unicast)
-        loop = loop or asyncio.get_event_loop()
+
-        with_unicast = loop.create_task(cls._discover_gateway(
+        done, pending = await asyncio.wait([
+            cls._discover_gateway(
                 lan_address, gateway_address, timeout, igd_args, loop, unicast=True
-        ))
+            ),
-        without_unicast = loop.create_task(cls._discover_gateway(
+            cls._discover_gateway(
                 lan_address, gateway_address, timeout, igd_args, loop, unicast=False
-        ))
+            )
-        await asyncio.wait([with_unicast, without_unicast], return_when=asyncio.tasks.FIRST_COMPLETED)
+        ], return_when=asyncio.tasks.FIRST_COMPLETED)
-        if with_unicast and not with_unicast.done():
+
-            with_unicast.cancel()
+        for task in pending:
-            if without_unicast.done():
+            task.cancel()
-                return without_unicast.result()
+        for task in done:
-        elif without_unicast and not without_unicast.done():
+            try:
-            without_unicast.cancel()
+                task.exception()
-            if with_unicast.done() and not with_unicast.cancelled():
+            except asyncio.CancelledError:
-                return with_unicast.result()
+                pass
-        else:
+
-            with_unicast.exception()
+        return list(done)[0].result()
-            without_unicast.exception()
-            return with_unicast.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)

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)