Use done and pending, add unit test
This commit is contained in:
parent
ee8ff02d52
commit
ca5ff5f225
2 changed files with 35 additions and 23 deletions
|
@ -155,8 +155,8 @@ class Gateway:
|
||||||
}
|
}
|
||||||
|
|
||||||
@classmethod
|
@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()
|
ignored: set = set()
|
||||||
required_commands = [
|
required_commands = [
|
||||||
'AddPortMapping',
|
'AddPortMapping',
|
||||||
|
@ -181,7 +181,7 @@ class Gateway:
|
||||||
required for required in required_commands if required not in gateway._registered_commands
|
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",
|
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)
|
ignored.add(datagram.location)
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
|
@ -197,26 +197,25 @@ class Gateway:
|
||||||
igd_args: OrderedDict = None, loop=None, unicast: bool = None):
|
igd_args: OrderedDict = None, loop=None, unicast: bool = None):
|
||||||
if unicast is not 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, unicast)
|
||||||
loop = loop or asyncio.get_event_loop()
|
|
||||||
with_unicast = loop.create_task(cls._discover_gateway(
|
done, pending = await asyncio.wait([
|
||||||
lan_address, gateway_address, timeout, igd_args, loop, unicast=True
|
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
|
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():
|
], return_when=asyncio.tasks.FIRST_COMPLETED)
|
||||||
with_unicast.cancel()
|
|
||||||
if without_unicast.done():
|
for task in pending:
|
||||||
return without_unicast.result()
|
task.cancel()
|
||||||
elif without_unicast and not without_unicast.done():
|
for task in done:
|
||||||
without_unicast.cancel()
|
try:
|
||||||
if with_unicast.done() and not with_unicast.cancelled():
|
task.exception()
|
||||||
return with_unicast.result()
|
except asyncio.CancelledError:
|
||||||
else:
|
pass
|
||||||
with_unicast.exception()
|
|
||||||
without_unicast.exception()
|
return list(done)[0].result()
|
||||||
return with_unicast.result()
|
|
||||||
|
|
||||||
async def discover_commands(self, loop=None):
|
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)
|
response, xml_bytes, get_err = await scpd_get(self.path.decode(), self.base_ip.decode(), self.port, loop=loop)
|
||||||
|
|
|
@ -1,3 +1,6 @@
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from aioupnp.fault import UPnPError
|
||||||
from tests import TestBase
|
from tests import TestBase
|
||||||
from tests.mocks import mock_tcp_and_udp
|
from tests.mocks import mock_tcp_and_udp
|
||||||
from collections import OrderedDict
|
from collections import OrderedDict
|
||||||
|
@ -62,6 +65,16 @@ class TestDiscoverDLinkDIR890L(TestBase):
|
||||||
'GetExternalIPAddress': 'urn:schemas-upnp-org:service:WANIPConnection:1'
|
'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):
|
async def test_discover_commands(self):
|
||||||
with mock_tcp_and_udp(self.loop, tcp_replies=self.replies):
|
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)
|
gateway = Gateway(self.reply, self.m_search_args, self.client_address, self.gateway_address)
|
||||||
|
|
Loading…
Reference in a new issue