Repository: sbidy/pywizlight

Issue:
Fix memory leak in discovery
There is a risk of leaking since `self.loop.call_later(1, self.broadcast_registration)` keeps calling and if `self.transport` is never cleared it will run as long as the object exists

Discussion:
sbidy: We should be aware of the #109. The IP can also be a DNS name. Especially in the HASS config flow if I enter a DNS name of a bulb manually.

https://github.com/sbidy/pywizlight/blob/cb6c00109fbbc18c7b0bf99466a0909fcaec1529/pywizlight/bulb.py#L301

I would suggest adding a `self.ip = socket.gethostbyname(ip)` to unify to IPs only.
bdraco: > @bdraco We should be aware of the #109. The IP can also be a DNS name. Especially in the HASS config flow if I enter a DNS name of a bulb manually.

Usually we don't want to accept dns names since if dns is having issues the device can't be setup.  As soon as discovery sees the bulb it will update it to an ip address via 

https://github.com/home-assistant/core/blob/dev/homeassistant/components/wiz/config_flow.py#L61

This takes care of the case where the bulb's dhcp reservation changes.



> ```python
> def _resolve_dnsname(self, dns_or_ip: str):
>         """Returns a IP in case of a DNS was given."""
>         try:
>             return socket.gethostbyname(dns_or_ip)
>         except socket.gaierror:
>             raise WizLightConnectionError("DNS name can not be resloved.")
> ```

If we do want to do allow a hostname, we need to use `loop.getaddrinfo()` since it is async safe.


Patch:
diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 0000000..5aeeb8f
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,5 @@
+[run]
+source = pywizlight
+
+omit =
+     pywizlight/tests/*
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index a727b64..75fd20e 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -14,7 +14,7 @@ permissions:
 
 jobs:
   run-linters:
-    name: Run linters
+    name: Run linters and tests
     runs-on: ubuntu-latest
 
     steps:
@@ -27,7 +27,7 @@ jobs:
           python-version: 3.8
 
       - name: Install Python dependencies
-        run: pip install black flake8 mypy pytest -e .
+        run: pip install black flake8 mypy pytest pytest-cov -e .
 
       - name: Run linters
         uses: samuelmeuli/lint-action@v1
@@ -40,3 +40,10 @@ jobs:
           flake8_args: "--exclude='test.py __init__.py' --max-line-length=120"
           auto_fix: false
           git_email: "sbidy@hotmail.com"
+
+      - name: Test with pytest
+        run: |
+            pytest --durations=0 --cov=pywizlight --cov-report term-missing --cov-report xml
+
+      - name: Upload codecov
+        uses: codecov/codecov-action@v2
diff --git a/pywizlight/bulb.py b/pywizlight/bulb.py
index ce572bf..7e46030 100755
--- a/pywizlight/bulb.py
+++ b/pywizlight/bulb.py
@@ -40,7 +40,7 @@ NEVER_TIME = -120.0
 _IGNORE_KEYS = {"src", "mqttCd", "ts", "rssi"}
 
 
-def states_match(old: dict[str, Any], new: dict[str, Any]) -> bool:
+def states_match(old: Dict[str, Any], new: Dict[str, Any]) -> bool:
     """Check if states match except for keys we do not want to callback on."""
     for key, val in new.items():
         if old.get(key) != val and key not in _IGNORE_KEYS:
@@ -285,6 +285,24 @@ class PilotParser:
             return None
 
 
+async def _send_udp_message_with_retry(
+    message: str,
+    transport: asyncio.DatagramTransport,
+    response_future: asyncio.Future,
+    ip: str,
+    port: int,
+) -> None:
+    """Send a UDP message multiple times until we reach the maximum or a response is recieved."""
+    data = message.encode("utf-8")
+    for send in range(MAX_SEND_DATAGRAMS):
+        if transport.is_closing() or response_future.done():
+            return
+        attempt = send + 1
+        _LOGGER.debug("%s: >> %s (%s/%s)", ip, data, attempt, MAX_SEND_DATAGRAMS)
+        transport.sendto(data, (ip, port))
+        await asyncio.sleep(SEND_INTERVAL)
+
+
 class wizlight:
     """Create an instance of a WiZ Light Bulb."""
 
@@ -333,7 +351,7 @@ class wizlight:
         if self.transport:
             return
         transport_proto = await self.loop.create_datagram_endpoint(
-            lambda: WizProtocol(on_response=self._on_response),
+            lambda: WizProtocol(on_response=self._on_response, on_error=self._on_error),
             remote_addr=(self.ip, self.port),
         )
         self.transport = cast(asyncio.DatagramTransport, transport_proto[0])
@@ -342,15 +360,17 @@ class wizlight:
     def register(self) -> None:
         """Call register to keep alive push updates."""
         if self.push_running:
-            asyncio.ensure_future(
+            asyncio.create_task(
                 self._async_send_register(PushManager().get().register_msg)
             )
             self.loop.call_later(PUSH_KEEP_ALIVE_INTERVAL, self.register)
 
     async def _async_send_register(self, message: str) -> None:
         """Send the registration message."""
-        with contextlib.suppress(WizLightTimeOutError):
+        try:
             await self.sendUDPMessage(message)
+        except (WizLightTimeOutError, WizLightConnectionError) as ex:
+            _LOGGER.debug("%s: Registration for push updates failed: %s", self.ip, ex)
 
     async def start_push(self, callback: Callable) -> None:
         """Start periodic register calls to get push updates via syncPilot."""
@@ -379,6 +399,11 @@ class wizlight:
         if self.response_future and not self.response_future.done():
             self.response_future.set_result(message)
 
+    def _on_error(self, exception: Optional[Exception]) -> None:
+        """Handle a protocol error."""
+        if exception and self.response_future and not self.response_future.done():
+            self.response_future.set_exception(exception)
+
     def _check_connection(self) -> None:
         """Check the connection to the bulb."""
         message = r'{"method":"getPilot","params":{}}'
@@ -556,40 +581,40 @@ class wizlight:
     async def sendUDPMessage(self, message: str) -> BulbResponse:
         """Send the UDP message to the bulb."""
         await self._ensure_connection()
-        data = message.encode("utf-8")
         assert self.transport is not None
         async with self.lock:
             self.response_future = asyncio.Future()
-            for send in range(MAX_SEND_DATAGRAMS):
-                attempt = send + 1
+            send_task = asyncio.create_task(
+                _send_udp_message_with_retry(
+                    message, self.transport, self.response_future, self.ip, self.port
+                )
+            )
+            try:
+                response = await asyncio.wait_for(
+                    asyncio.shield(self.response_future), timeout=TIMEOUT
+                )
+            except OSError as ex:
+                raise WizLightConnectionError(str(ex)) from ex
+            except asyncio.TimeoutError:
                 _LOGGER.debug(
-                    "%s: >> %s (%s/%s)", self.ip, data, attempt, MAX_SEND_DATAGRAMS
+                    "%s: Timed out waiting for response to %s after %s tries.",
+                    self.ip,
+                    message,
+                    MAX_SEND_DATAGRAMS,
                 )
-                self.transport.sendto(data, (self.ip, self.port))
-                try:
-                    response = await asyncio.wait_for(
-                        asyncio.shield(self.response_future), timeout=SEND_INTERVAL
-                    )
-                except asyncio.TimeoutError:
-                    _LOGGER.debug(
-                        "%s: Timed out waiting for response to %s (%s/%s)",
-                        self.ip,
-                        message,
-                        attempt,
-                        MAX_SEND_DATAGRAMS,
-                    )
-                    if attempt == MAX_SEND_DATAGRAMS:
-                        raise WizLightTimeOutError("The request to the bulb timed out")
-                else:
-                    break
-
+                raise WizLightTimeOutError("The request to the bulb timed out")
+            finally:
+                await asyncio.sleep(0)  # Try to let the task finish without cancelation
+                if not send_task.done():
+                    send_task.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await send_task
         resp = json.loads(response.decode())
-        if "error" not in resp:
-            return resp
-        elif resp["error"]["code"] == -32601:
-            raise WizLightMethodNotFound("Method not found; maybe older bulb FW?")
-        else:
+        if "error" in resp:
+            if resp["error"]["code"] == -32601:
+                raise WizLightMethodNotFound("Method not found; maybe older bulb FW?")
             raise WizLightConnectionError(f'Error recieved: {resp["error"]}')
+        return resp
 
     async def async_close(self):
         """Close the transport."""
@@ -606,5 +631,5 @@ class wizlight:
 
     def __del__(self):
         """Close the connection when the object is destroyed."""
-        if self.transport:
+        if self.transport and not self.transport.is_closing():
             self.loop.call_soon_threadsafe(self._async_close)
diff --git a/pywizlight/discovery.py b/pywizlight/discovery.py
index b72135c..52d46df 100644
--- a/pywizlight/discovery.py
+++ b/pywizlight/discovery.py
@@ -3,14 +3,22 @@ import asyncio
 import dataclasses
 import json
 import logging
-import socket
 from asyncio import DatagramTransport, BaseTransport, AbstractEventLoop, Future
-from typing import TYPE_CHECKING, Dict, List, Tuple, Optional, Any
+from typing import cast, Dict, List, Tuple, Optional, Any
 
 from pywizlight import wizlight
+from pywizlight.utils import create_udp_broadcast_socket
 
 _LOGGER = logging.getLogger(__name__)
 
+PORT = 38899
+DEFAULT_WAIT_TIME = 5.0
+
+# Note: The IP and address we give the bulb does not matter because
+# we have register set to false which is telling the bulb to remove
+# the registration
+REGISTER_MESSAGE = b'{"method":"registration","params":{"phoneMac":"AAAAAAAAAAAA","register":false,"phoneIp":"1.2.3.4","id":"1"}}'  # noqa: E501
+
 
 @dataclasses.dataclass(frozen=True)
 class DiscoveredBulb:
@@ -19,18 +27,6 @@ class DiscoveredBulb:
     ip_address: str
     mac_address: str
 
-    @staticmethod
-    def create_bulb_from_message(
-        raw_addr: Tuple[str, int], announce_message: str
-    ) -> "DiscoveredBulb":
-        """Create announce message."""
-        ip_address = raw_addr[0]
-        _LOGGER.debug(f"Found bulb with IP: {ip_address}")
-        return DiscoveredBulb(
-            ip_address=ip_address,
-            mac_address=json.loads(announce_message)["result"]["mac"],
-        )
-
 
 @dataclasses.dataclass
 class BulbRegistry:
@@ -66,41 +62,28 @@ class BroadcastProtocol(asyncio.DatagramProtocol):
 
     def connection_made(self, transport: BaseTransport) -> None:
         """Init connection to socket and register broadcasts."""
-        if TYPE_CHECKING:
-            # TODO: if this is made a runtime check, it needs some changes (see #94)
-            assert isinstance(
-                transport, BaseTransport
-            )  # Required to keep this liskov-safe
-        self.transport = transport  # type: ignore
-        sock = transport.get_extra_info("socket")
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+        self.transport = cast(DatagramTransport, transport)
         self.broadcast_registration()
 
     def broadcast_registration(self) -> None:
         """Send a registration method as UDP broadcast."""
-        # Note: The IP and address we give the bulb here don't seem to matter for our
-        # intents and purposes, so they're hardcoded to technically valid dummy data.
-        # Fix for async problems if broaddcast_registration is called twice! See #13.
-        # dirty dirty hack
         if not self.transport:
             return
-        try:
-            register_method = r'{"method":"registration","params":{"phoneMac":"AAAAAAAAAAAA","register":false,"phoneIp":"1.2.3.4","id":"1"}}'  # noqa: E501
-            self.transport.sendto(
-                register_method.encode(), (self.broadcast_address, 38899)
-            )
-            self.loop.call_later(1, self.broadcast_registration)
-        except AttributeError:
-            pass
+        self.transport.sendto(REGISTER_MESSAGE, (self.broadcast_address, 38899))
+        self.loop.call_later(1, self.broadcast_registration)
 
     def datagram_received(self, data: bytes, addr: Tuple[str, int]) -> None:
         """Receive data from broadcast."""
         _LOGGER.debug(f"Received data {data!r} from IP address {addr}")
-        decoded_message = data.decode()
-        if """"success":true""" in decoded_message:
-            bulb = DiscoveredBulb.create_bulb_from_message(addr, decoded_message)
-            self.registry.register(bulb)
+        try:
+            resp = json.loads(data.decode())
+        except json.JSONDecodeError:
+            _LOGGER.error("%s: Sent invalid message: %s", addr, data)
+            return
+        mac = resp.get("result", {}).get("mac")
+        if mac:
+            _LOGGER.debug("Found bulb with IP: %s and MAC: %s", addr[0], mac)
+            self.registry.register(DiscoveredBulb(addr[0], mac))
 
     def connection_lost(self, exc: Any) -> None:
         """Return connection error."""
@@ -109,10 +92,11 @@ class BroadcastProtocol(asyncio.DatagramProtocol):
             self.future.set_result(None)
         else:
             self.future.set_exception(exc)
+        self.transport = None
 
 
 async def find_wizlights(
-    wait_time: float = 5, broadcast_address: str = "255.255.255.255"
+    wait_time: float = DEFAULT_WAIT_TIME, broadcast_address: str = "255.255.255.255"
 ) -> List[DiscoveredBulb]:
     """Start discovery and return list of IPs of the bulbs."""
     registry = BulbRegistry()
@@ -120,27 +104,24 @@ async def find_wizlights(
     future = loop.create_future()
     transport, protocol = await loop.create_datagram_endpoint(
         lambda: BroadcastProtocol(loop, registry, broadcast_address, future),
-        local_addr=("0.0.0.0", 38899),
+        sock=create_udp_broadcast_socket(PORT),
     )
-    try:
-        await asyncio.sleep(wait_time)
-    finally:
-        transport.close()
-        await future
-        bulbs = registry.bulbs()
-        for bulb in bulbs:
-            _LOGGER.info(
-                f"Discovered bulb {bulb.ip_address} with MAC {bulb.mac_address}"
-            )
-        return bulbs
-
-
-async def discover_lights(broadcast_space: str = "255.255.255.255") -> List[wizlight]:
-    """Find lights and return list with wizlight objects."""
-    discovered_IPs = await find_wizlights(broadcast_address=broadcast_space)
-    # empty list for adding bulbs
-    bulbs = []
-    # create light entities from register
-    for entries in discovered_IPs:
-        bulbs.append(wizlight(ip=entries.ip_address, mac=entries.mac_address))
+    await asyncio.sleep(wait_time)
+    transport.close()
+    await future
+    bulbs = registry.bulbs()
+    for bulb in bulbs:
+        _LOGGER.info(f"Discovered bulb {bulb.ip_address} with MAC {bulb.mac_address}")
     return bulbs
+
+
+async def discover_lights(
+    broadcast_space: str = "255.255.255.255", wait_time: float = DEFAULT_WAIT_TIME
+) -> List[wizlight]:
+    """Find lights and return list with wizlight objects."""
+    discovered_IPs = await find_wizlights(
+        wait_time=wait_time, broadcast_address=broadcast_space
+    )
+    return [
+        wizlight(ip=entry.ip_address, mac=entry.mac_address) for entry in discovered_IPs
+    ]
diff --git a/pywizlight/protocol.py b/pywizlight/protocol.py
index 1c3db34..05d9a37 100755
--- a/pywizlight/protocol.py
+++ b/pywizlight/protocol.py
@@ -10,10 +10,12 @@ class WizProtocol(asyncio.DatagramProtocol):
     def __init__(
         self,
         on_response: Callable[[bytes, Tuple[str, int]], None],
+        on_error: Optional[Callable[[Optional[Exception]], None]] = None,
     ) -> None:
         """Init the protocol."""
         self.transport = None
         self.on_response = on_response
+        self.on_error = on_error
 
     def datagram_received(self, data: bytes, addr: Tuple[str, int]) -> None:
         """Trigger on_response."""
@@ -22,6 +24,8 @@ class WizProtocol(asyncio.DatagramProtocol):
     def error_received(self, ex: Optional[Exception]) -> None:
         """Handle error."""
         _LOGGER.debug("WizProtocol error: %s", ex)
+        if self.on_error is not None:
+            self.on_error(ex)
 
     def connection_lost(self, ex: Optional[Exception]) -> None:
         """The connection is lost."""
diff --git a/pywizlight/push_manager.py b/pywizlight/push_manager.py
index ff9b9b4..268b63d 100755
--- a/pywizlight/push_manager.py
+++ b/pywizlight/push_manager.py
@@ -7,7 +7,7 @@ import socket
 from typing import Callable, Dict, Optional, Tuple, cast
 
 from pywizlight.protocol import WizProtocol
-from pywizlight.utils import to_wiz_json
+from pywizlight.utils import to_wiz_json, create_udp_socket
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -15,14 +15,6 @@ RESPOND_PORT = 38899
 LISTEN_PORT = 38900
 
 
-def create_udp_socket(listen_port: int) -> socket.socket:
-    """Create a udp socket used for communicating with the device."""
-    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    sock.bind(("", listen_port))
-    sock.setblocking(False)
-    return sock
-
-
 def generate_mac():
     """Generates a fake mac address."""
     return "{}{}{}{}{}{}{}{}{}{}{}{}".format(
@@ -65,7 +57,7 @@ class PushManager:
         self.push_running = False
         self.lock = asyncio.Lock()
         self.loop = asyncio.get_event_loop()
-        self.subscriptions: Dict[str, Callable] = {}
+        self.subscriptions: Dict[str, Callable[[Dict, Tuple[str, int]], None]] = {}
         self.register_msg: Optional[str] = None
 
     async def start(self, target_ip: str) -> bool:
@@ -125,7 +117,7 @@ class PushManager:
 
         def _cancel():
             del self.subscriptions[mac]
-            asyncio.ensure_future(self.stop_if_no_subs())
+            asyncio.create_task(self.stop_if_no_subs())
 
         return _cancel
 
@@ -139,12 +131,9 @@ class PushManager:
             return
         method = resp.get("method")
         mac = resp.get("params", {}).get("mac")
-        if method != "syncPilot":
-            return
-        if self.push_transport:
-            data = to_wiz_json({"method": "syncPilot", "result": {"mac": mac}}).encode()
-            _LOGGER.debug("%s: PUSH ACK >> %s", (addr[0], RESPOND_PORT), data)
-            self.push_transport.sendto(data, (addr[0], RESPOND_PORT))
-        subscription = self.subscriptions.get(mac)
-        if subscription:
-            subscription(resp, addr)
+        # We used to send an response to `syncPilot` messages
+        # but the devices keeps sending them even if we do not
+        # so we no longer send them since all it effectively
+        # does it generate additional network traffic.
+        if method == "syncPilot" and mac in self.subscriptions:
+            self.subscriptions[mac](resp, addr)
diff --git a/pywizlight/utils.py b/pywizlight/utils.py
index 14111c2..b1ab2fd 100644
--- a/pywizlight/utils.py
+++ b/pywizlight/utils.py
@@ -1,4 +1,5 @@
 import json
+import socket
 
 
 def hex_to_percent(hex_value: float) -> float:
@@ -14,3 +15,21 @@ def percent_to_hex(percent: float) -> int:
 def to_wiz_json(dump_obj):
     """Convert an object to wiz json."""
     return json.dumps(dump_obj, separators=(",", ":"))
+
+
+def create_udp_broadcast_socket(listen_port: int) -> socket.socket:
+    """Create a udp broadcast socket used for communicating with the device."""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+    sock.bind(("", listen_port))
+    sock.setblocking(False)
+    return sock
+
+
+def create_udp_socket(listen_port: int) -> socket.socket:
+    """Create a udp socket used for communicating with the device."""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    sock.bind(("", listen_port))
+    sock.setblocking(False)
+    return sock


Test patch:
diff --git a/pywizlight/tests/fake_bulb.py b/pywizlight/tests/fake_bulb.py
index d07d3aa..f7acacf 100644
--- a/pywizlight/tests/fake_bulb.py
+++ b/pywizlight/tests/fake_bulb.py
@@ -85,7 +85,7 @@ class BulbUDPRequestHandlerBase(socketserver.DatagramRequestHandler):
         if method == "setPilot":
             return_data = self.setPilot(json_data)
             self.wfile.write(return_data)
-        if method == "getPilot":
+        elif method == "getPilot":
             print(f"Response:{json.dumps(self.pilot_state)!r}")
             self.wfile.write(bytes(json.dumps(self.pilot_state), "utf-8"))
         elif method == "getSystemConfig":
diff --git a/pywizlight/tests/test_bulb.py b/pywizlight/tests/test_bulb.py
index cd11ae2..0187b74 100644
--- a/pywizlight/tests/test_bulb.py
+++ b/pywizlight/tests/test_bulb.py
@@ -35,9 +35,15 @@ async def bad_bulb() -> AsyncGenerator[wizlight, None]:
 @pytest.mark.asyncio
 async def test_Bulb_Discovery(correct_bulb: wizlight) -> None:
     """Test discovery function."""
-    bulbs = await discover_lights(broadcast_space="192.168.178.255")
+    # Use a random available port since otherwise the
+    # test may fail
+    with patch("pywizlight.discovery.PORT", 0):
+        bulbs = await discover_lights(broadcast_space="192.168.178.255", wait_time=0.02)
     for bulb in bulbs:
-        state = await bulb.updateState()
+        with patch("pywizlight.bulb.SEND_INTERVAL", 0.01), patch(
+            "pywizlight.bulb.TIMEOUT", 0.01
+        ):
+            state = await bulb.updateState()
         assert state and state.get_state() is False