From ee8ff02d52de88667a1e3f07ffcf82347a62ac92 Mon Sep 17 00:00:00 2001
From: hackrush <atulshyan96@gmail.com>
Date: Sat, 2 Feb 2019 12:22:20 +0530
Subject: [PATCH] Call `.exception()` if both discoveries fail together, else
 return result

---
 aioupnp/gateway.py | 33 ++++++++++++++++++++-------------
 aioupnp/upnp.py    |  2 +-
 tests/test_cli.py  |  2 +-
 3 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/aioupnp/gateway.py b/aioupnp/gateway.py
index bdd918f..0749427 100644
--- a/aioupnp/gateway.py
+++ b/aioupnp/gateway.py
@@ -197,19 +197,26 @@ 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)
-            return await cls._discover_gateway(lan_address, gateway_address, timeout, igd_args, loop)
-        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):
-            task.cancel()
-        result = list(done)[0].result()
-        return result
+        loop = loop or asyncio.get_event_loop()
+        with_unicast = loop.create_task(cls._discover_gateway(
+            lan_address, gateway_address, timeout, igd_args, loop, unicast=True
+        ))
+        without_unicast = loop.create_task(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)
+        if with_unicast and not with_unicast.done():
+            with_unicast.cancel()
+            if without_unicast.done():
+                return without_unicast.result()
+        elif without_unicast and not without_unicast.done():
+            without_unicast.cancel()
+            if with_unicast.done() and not with_unicast.cancelled():
+                return with_unicast.result()
+        else:
+            with_unicast.exception()
+            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)
diff --git a/aioupnp/upnp.py b/aioupnp/upnp.py
index e184a36..9eab4e9 100644
--- a/aioupnp/upnp.py
+++ b/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)):
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 91def3e..e1c4da5 100644
--- a/tests/test_cli.py
+++ b/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(