diff --git a/aioupnp/gateway.py b/aioupnp/gateway.py index d772d98..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: @@ -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() - return result + 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/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( 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)