From daaa41ba3daa8acb40dafbef84a9b9b61134d031 Mon Sep 17 00:00:00 2001 From: hackrush Date: Sat, 2 Feb 2019 12:11:02 +0530 Subject: [PATCH 1/3] pass `unicast` if not None --- aioupnp/gateway.py | 1 + 1 file changed, 1 insertion(+) diff --git a/aioupnp/gateway.py b/aioupnp/gateway.py index d772d98..bdd918f 100644 --- a/aioupnp/gateway.py +++ b/aioupnp/gateway.py @@ -196,6 +196,7 @@ 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, unicast) return await cls._discover_gateway(lan_address, gateway_address, timeout, igd_args, loop) done, pending = await asyncio.wait([ cls._discover_gateway( -- 2.45.2 From ee8ff02d52de88667a1e3f07ffcf82347a62ac92 Mon Sep 17 00:00:00 2001 From: hackrush Date: Sat, 2 Feb 2019 12:22:20 +0530 Subject: [PATCH 2/3] 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( -- 2.45.2 From ca5ff5f2257f79557e534e0ae482131d2fef9462 Mon Sep 17 00:00:00 2001 From: hackrush Date: Wed, 6 Feb 2019 00:46:49 +0530 Subject: [PATCH 3/3] Use done and pending, add unit test --- aioupnp/gateway.py | 45 +++++++++++++++++++++---------------------- tests/test_gateway.py | 13 +++++++++++++ 2 files changed, 35 insertions(+), 23 deletions(-) diff --git a/aioupnp/gateway.py b/aioupnp/gateway.py index 0749427..ad04143 100644 --- a/aioupnp/gateway.py +++ b/aioupnp/gateway.py @@ -155,8 +155,8 @@ class Gateway: } @classmethod - async def _discover_gateway(cls, lan_address: str, gateway_address: str, timeout: int=30, - igd_args: OrderedDict=None, loop=None, unicast: bool=False): + async def _discover_gateway(cls, lan_address: str, gateway_address: str, timeout: int = 30, + igd_args: OrderedDict = None, loop=None, unicast: bool = False): ignored: set = set() required_commands = [ 'AddPortMapping', @@ -181,7 +181,7 @@ class Gateway: required for required in required_commands if required not in gateway._registered_commands ] log.debug("found gateway %s at %s, but it does not implement required soap commands: %s", - gateway.manufacturer_string, gateway.location, not_met) + gateway.manufacturer_string, gateway.location, not_met) ignored.add(datagram.location) continue else: @@ -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( - 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() + + 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 pending: + task.cancel() + for task in done: + 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) diff --git a/tests/test_gateway.py b/tests/test_gateway.py index f232659..c105f11 100644 --- a/tests/test_gateway.py +++ b/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) -- 2.45.2