Repository: postlund/pyatv

Issue:
Zeroconf response incomplete due to response not fitting in MTU
<!--
READ THIS FIRST!

Thank you for reporting any bugs you find! Before opening a new issue, please
look through already open issues (someone else might have reported it already).
Also, take a look in the FAQ and especially the *Known Issues* section at the
bottom. Here is a direct link:

https://pyatv.dev/support/faq/
-->

**Describe the bug**
Original issue was discovered in https://github.com/home-assistant/core/issues/52862#issuecomment-885229784

The problem is that now when pyatv asks for a lot of services in the same request, the anticipated response might not fit in the packet MTU (roughly 1450 bytes or so). This makes the sender ignore records, thus making the final response incomplete. In this case everything was present, but only a AAAA record no A record. So I have to split and make several requests, see mentioned issue above for details.

**To Reproduce**
Scan (regular or unicast) for a device producing too large responses. Device is not found.

**Expected behavior**
Device is found

**System Setup (please complete the following information):**
 - OS: any
 - Python: any
 - pyatv: 0.8.0-0.8.1
 - Apple TV: any

**Additional context**
Implicitly introduced in 0.8.0 as many additional services were added.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl) Report
> Merging [#1235](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl) (5856097) into [master](https://codecov.io/gh/postlund/pyatv/commit/e1dab5721a0d9ea589eb1f42da1f051b1cb174d9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl) (e1dab57) will **decrease** coverage by `0.37%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/postlund/pyatv/pull/1235/graphs/tree.svg?width=650&height=150&src=pr&token=wb6azJjRSY&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl)](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl)

```diff
@@            Coverage Diff             @@
##           master    #1235      +/-   ##
==========================================
- Coverage   95.25%   94.87%   -0.38%     
==========================================
  Files          58       58              
  Lines        6148     6191      +43     
==========================================
+ Hits         5856     5874      +18     
- Misses        292      317      +25     
```


| [Impacted Files](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl) | Coverage Δ | |
|---|---|---|
| [pyatv/support/mdns.py](https://codecov.io/gh/postlund/pyatv/pull/1235/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl#diff-cHlhdHYvc3VwcG9ydC9tZG5zLnB5) | `89.73% <100.00%> (-0.03%)` | :arrow_down: |
| [pyatv/raop/raop.py](https://codecov.io/gh/postlund/pyatv/pull/1235/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl#diff-cHlhdHYvcmFvcC9yYW9wLnB5) | `85.71% <0.00%> (-3.93%)` | :arrow_down: |
| [pyatv/raop/audio\_source.py](https://codecov.io/gh/postlund/pyatv/pull/1235/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl#diff-cHlhdHYvcmFvcC9hdWRpb19zb3VyY2UucHk=) | `86.87% <0.00%> (-2.50%)` | :arrow_down: |
| [pyatv/mrp/connection.py](https://codecov.io/gh/postlund/pyatv/pull/1235/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl#diff-cHlhdHYvbXJwL2Nvbm5lY3Rpb24ucHk=) | `81.25% <0.00%> (-2.09%)` | :arrow_down: |
| [pyatv/mrp/protocol.py](https://codecov.io/gh/postlund/pyatv/pull/1235/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl#diff-cHlhdHYvbXJwL3Byb3RvY29sLnB5) | `93.70% <0.00%> (-1.40%)` | :arrow_down: |
| [pyatv/companion/\_\_init\_\_.py](https://codecov.io/gh/postlund/pyatv/pull/1235/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl#diff-cHlhdHYvY29tcGFuaW9uL19faW5pdF9fLnB5) | `91.07% <0.00%> (+0.79%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl). Last update [e1dab57...5856097](https://codecov.io/gh/postlund/pyatv/pull/1235?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Pierre+St%C3%A5hl).

postlund: Everything implemented and verified in https://github.com/home-assistant/core/issues/52862

Patch:
diff --git a/pyatv/support/mdns.py b/pyatv/support/mdns.py
index 20c662e6..6e37ff8a 100644
--- a/pyatv/support/mdns.py
+++ b/pyatv/support/mdns.py
@@ -4,6 +4,7 @@ from ipaddress import IPv4Address, ip_address
 import logging
 import math
 import socket
+from types import SimpleNamespace
 import typing
 import weakref
 
@@ -22,6 +23,11 @@ from pyatv.support.dns import (
 
 _LOGGER = logging.getLogger(__name__)
 
+# Number of services to include in each request
+SERVICES_PER_MSG = 3
+
+SLEEP_PROXY_SERVICE = "_sleep-proxy._udp.local"
+
 # This module produces a lot of debug output, use a dedicated log level.
 # Maybe move this to top-level support later?
 TRAFFIC_LEVEL = logging.DEBUG - 5
@@ -67,13 +73,20 @@ def _decode_properties(
     return CaseInsensitiveDict({k: _decode(v) for k, v in properties.items()})
 
 
-def create_request(
-    services: typing.List[str], qtype: QueryType = QueryType.PTR
-) -> bytes:
-    """Create a new DnsMessage requesting specified services."""
-    msg = DnsMessage(0x35FF)
-    msg.questions += [DnsQuestion(s, qtype, 0x8001) for s in services]
-    return msg.pack()
+def create_service_queries(
+    services: typing.List[str], qtype: QueryType
+) -> typing.List[bytes]:
+    """Create service request messages."""
+    queries: typing.List[bytes] = []
+    for i in range(math.ceil(len(services) / SERVICES_PER_MSG)):
+        service_chunk = services[i * SERVICES_PER_MSG : i * SERVICES_PER_MSG + 4]
+
+        msg = DnsMessage(0x35FF)
+        msg.questions += [DnsQuestion(s, qtype, 0x8001) for s in service_chunk]
+        msg.questions += [DnsQuestion(SLEEP_PROXY_SERVICE, qtype, 0x8001)]
+
+        queries.append(msg.pack())
+    return queries
 
 
 def _get_model(services: typing.List[Service]) -> typing.Optional[str]:
@@ -87,57 +100,83 @@ def _first_rd(qtype: QueryType, entries: typing.Dict[int, typing.List[DnsResourc
     return entries[qtype][0].rd if qtype in entries else None
 
 
-def parse_services(  # pylint: disable=too-many-locals
-    message: DnsMessage,
-) -> typing.List[Service]:
-    """Parse DNS response into Service objects."""
-    table: typing.Dict[str, typing.Dict[int, typing.List[DnsResource]]] = {}
-    ptrs: typing.Dict[str, str] = {}  # qname -> real name
-    results: typing.Dict[str, Service] = {}
+class ServiceParser:
+    """Parse zeroconf services from records in DNS messages."""
 
-    # Create a global table with all records
-    for record in message.answers + message.resources:
-        if record.qtype == QueryType.PTR and record.qname.startswith("_"):
-            ptrs[record.qname] = record.rd
-        else:
-            entry = table.setdefault(record.qname, {})
-            if record.qtype not in entry:
-                entry[record.qtype] = []
-            entry[record.qtype].append(record)
+    def __init__(self) -> None:
+        """Initialize a new ServiceParser instance."""
+        self.table: typing.Dict[str, typing.Dict[int, typing.List[DnsResource]]] = {}
+        self.ptrs: typing.Dict[str, str] = {}  # qname -> real name
+        self._cache: typing.Optional[typing.List[Service]] = None
 
-    # Build services
-    for service, device in table.items():
-        try:
-            service_name = ServiceInstanceName.split_name(service)
-        except ValueError:
-            continue
-
-        srv_rd = _first_rd(QueryType.SRV, device)
-        target = srv_rd["target"] if srv_rd else None
-
-        target_records = table.get(typing.cast(str, target), {}).get(QueryType.A, [])
-        address = None
-
-        # Pick one address that is not link-local
-        for addr in [IPv4Address(record.rd) for record in target_records]:
-            if not addr.is_link_local:
-                address = addr
-                break
-
-        results[service] = Service(
-            service_name.ptr_name,
-            typing.cast(str, service_name.instance),
-            address,
-            srv_rd["port"] if srv_rd else 0,
-            _decode_properties(_first_rd(QueryType.TXT, device) or {}),
-        )
+    def add_message(self, message: DnsMessage) -> "ServiceParser":
+        """Add message to with records to parse."""
+        self._cache = None
 
-    # If there are PTRs to unknown services, create placeholders
-    for qname, real_name in ptrs.items():
-        if real_name not in results:
-            results[real_name] = Service(qname, real_name.split(".")[0], None, 0, {})
+        for record in message.answers + message.resources:
+            if record.qtype == QueryType.PTR and record.qname.startswith("_"):
+                self.ptrs[record.qname] = record.rd
+            else:
+                entry = self.table.setdefault(record.qname, {})
+                if record.qtype not in entry:
+                    entry[record.qtype] = []
+
+                if record not in entry[record.qtype]:
+                    entry[record.qtype].append(record)
+        return self
 
-    return list(results.values())
+    def parse(self) -> typing.List[Service]:
+        """Parse records and return services."""
+        if self._cache:
+            return self._cache
+
+        results: typing.Dict[str, Service] = {}
+
+        # Build services
+        for service, device in self.table.items():
+            try:
+                service_name = ServiceInstanceName.split_name(service)
+            except ValueError:
+                continue
+
+            srv_rd = _first_rd(QueryType.SRV, device)
+            target = srv_rd["target"] if srv_rd else None
+
+            target_records = self.table.get(typing.cast(str, target), {}).get(
+                QueryType.A, []
+            )
+            address = None
+
+            # Pick one address that is not link-local
+            for addr in [IPv4Address(record.rd) for record in target_records]:
+                if not addr.is_link_local:
+                    address = addr
+                    break
+
+            results[service] = Service(
+                service_name.ptr_name,
+                typing.cast(str, service_name.instance),
+                address,
+                srv_rd["port"] if srv_rd else 0,
+                _decode_properties(_first_rd(QueryType.TXT, device) or {}),
+            )
+
+        # If there are PTRs to unknown services, create placeholders
+        for qname, real_name in self.ptrs.items():
+            if real_name not in results:
+                results[real_name] = Service(
+                    qname, real_name.split(".")[0], None, 0, {}
+                )
+        self._cache = list(results.values())
+        return self._cache
+
+
+class QueryResponse(SimpleNamespace):
+    """Hold DNS query response records."""
+
+    count: int
+    deep_sleep: bool
+    parser: ServiceParser
 
 
 class UnicastDnsSdClientProtocol(asyncio.Protocol):
@@ -145,12 +184,13 @@ class UnicastDnsSdClientProtocol(asyncio.Protocol):
 
     def __init__(self, services: typing.List[str], host: str, timeout: int):
         """Initialize a new UnicastDnsSdClientProtocol."""
-        self.message = create_request(services)
+        self.queries = create_service_queries(services, QueryType.PTR)
         self.host = host
         self.timeout = timeout
         self.transport = None
+        self.parser: ServiceParser = ServiceParser()
         self.semaphore: asyncio.Semaphore = asyncio.Semaphore(value=0)
-        self.result: DnsMessage = DnsMessage()
+        self.received_responses: int = 0
         self._task: typing.Optional[asyncio.Future] = None
 
     async def get_response(self) -> Response:
@@ -167,7 +207,8 @@ class UnicastDnsSdClientProtocol(asyncio.Protocol):
                 await self._task
             except asyncio.CancelledError:
                 pass
-        services = parse_services(self.result)
+
+        services = self.parser.parse()
         return Response(
             services=services,
             deep_sleep=False,
@@ -181,14 +222,15 @@ class UnicastDnsSdClientProtocol(asyncio.Protocol):
 
     async def _resend_loop(self):
         for _ in range(math.ceil(self.timeout)):
-            log_binary(
-                _LOGGER,
-                "Sending DNS request to " + self.host,
-                level=TRAFFIC_LEVEL,
-                Data=self.message,
-            )
+            for query in self.queries:
+                log_binary(
+                    _LOGGER,
+                    "Sending DNS request to " + self.host,
+                    level=TRAFFIC_LEVEL,
+                    Data=query,
+                )
 
-            self.transport.sendto(self.message)
+                self.transport.sendto(query)
             await asyncio.sleep(1)
 
     def datagram_received(self, data: bytes, _) -> None:
@@ -198,12 +240,17 @@ class UnicastDnsSdClientProtocol(asyncio.Protocol):
             "Received DNS response from " + self.host,
             level=TRAFFIC_LEVEL,
             Data=data,
+            Index=self.received_responses + 1,
+            Total=len(self.queries),
         )
 
-        self.result = DnsMessage().unpack(data)
-        self._finished()
-        if self.transport:
-            self.transport.close()
+        self.parser.add_message(DnsMessage().unpack(data))
+        self.received_responses += 1
+
+        if self.received_responses == len(self.queries):
+            self._finished()
+            if self.transport:
+                self.transport.close()
 
     def error_received(self, exc) -> None:
         """Error received during communication."""
@@ -285,13 +332,14 @@ class MulticastDnsSdClientProtocol:  # pylint: disable=too-many-instance-attribu
         """Initialize a new MulticastDnsSdClientProtocol."""
         self.loop = loop
         self.services = services
-        self.message = create_request(services)
+        self.queries = create_service_queries(services, QueryType.PTR)
+        self.query_responses: typing.Dict[str, QueryResponse] = {}
         self.address = address
         self.port = port
         self.end_condition = end_condition or (lambda _: False)
         self.semaphore: asyncio.Semaphore = asyncio.Semaphore(value=0)
-        self.responses: typing.Dict[IPv4Address, Response] = {}
-        self._unicasts: typing.Dict[IPv4Address, bytes] = {}
+        self.parser = ServiceParser()
+        self._unicasts: typing.Dict[IPv4Address, typing.List[bytes]] = {}
         self._task: typing.Optional[asyncio.Future] = None
         self._receivers: typing.List[asyncio.BaseProtocol] = []
 
@@ -304,7 +352,7 @@ class MulticastDnsSdClientProtocol:  # pylint: disable=too-many-instance-attribu
 
         self._receivers.append(protocol)
 
-    async def get_response(self, timeout: int) -> typing.Dict[IPv4Address, Response]:
+    async def get_response(self, timeout: int) -> typing.List[Response]:
         """Get respoonse with a maximum timeout."""
         # Semaphore used here as a quick-bailout when testing
         try:
@@ -320,28 +368,39 @@ class MulticastDnsSdClientProtocol:  # pylint: disable=too-many-instance-attribu
                 except asyncio.CancelledError:
                     pass
                 self._task = None
-        return self.responses
 
-    async def _resend_loop(self, timeout):
-        for _ in range(math.ceil(timeout)):
-            log_binary(
-                _LOGGER,
-                f"Sending multicast DNS request to {self.address}:{self.port}",
-                level=TRAFFIC_LEVEL,
-                Data=self.message,
+        def _to_response(query_response: QueryResponse):
+            services = query_response.parser.parse()
+            return Response(
+                services=services,
+                deep_sleep=query_response.deep_sleep,
+                model=_get_model(services),
             )
 
-            self._sendto(self.message, (self.address, self.port))
+        return [_to_response(response) for response in self.query_responses.values()]
 
-            # Send unicast requests if devices are sleeping
-            for address, message in self._unicasts.items():
+    async def _resend_loop(self, timeout):
+        for _ in range(math.ceil(timeout)):
+            for query in self.queries:
                 log_binary(
                     _LOGGER,
-                    f"Sending unicast DNS request to {address}:{self.port}",
+                    f"Sending multicast DNS request to {self.address}:{self.port}",
                     level=TRAFFIC_LEVEL,
-                    Data=message,
+                    Data=query,
                 )
-                self._sendto(message, (address, self.port))
+
+                self._sendto(query, (self.address, self.port))
+
+            # Send unicast requests if devices are sleeping
+            for address, queries in self._unicasts.items():
+                for query in queries:
+                    log_binary(
+                        _LOGGER,
+                        f"Sending unicast DNS request to {address}:{self.port}",
+                        level=TRAFFIC_LEVEL,
+                        Data=query,
+                    )
+                    self._sendto(query, (address, self.port))
 
             await asyncio.sleep(1)
 
@@ -361,41 +420,53 @@ class MulticastDnsSdClientProtocol:  # pylint: disable=too-many-instance-attribu
             Data=data,
         )
 
+        query_resp = self.query_responses.setdefault(
+            addr[0], QueryResponse(count=0, deep_sleep=False, parser=ServiceParser())
+        )
+
         # Suppress decode errors for now (but still log)
         try:
-            services = parse_services(DnsMessage().unpack(data))
+            decoded_msg = DnsMessage().unpack(data)
+
+            parser = ServiceParser()
+            services = parser.add_message(decoded_msg).parse()
         except UnicodeDecodeError:
             log_binary(_LOGGER, "Failed to decode message", Msg=data)
             return
+        else:
+            if not services:
+                return
 
         # Ignore responses from other services
         for service in services:
-            if (
-                service.type not in self.services
-                and service.type != DEVICE_INFO_SERVICE
+            if not (
+                service.type in self.services
+                or service.type in [DEVICE_INFO_SERVICE, SLEEP_PROXY_SERVICE]
             ):
                 return
 
         is_sleep_proxy = all(service.port == 0 for service in services)
+        query_resp.count += 1
+        query_resp.deep_sleep |= is_sleep_proxy
+        query_resp.parser.add_message(decoded_msg)
+
         if is_sleep_proxy:
-            self._unicasts[addr[0]] = create_request(
+            self._unicasts[addr[0]] = create_service_queries(
                 [service.name + "." + service.type for service in services],
-                qtype=QueryType.ANY,
+                QueryType.ANY,
             )
-        else:
+        elif query_resp.count >= len(self.queries):
             response = Response(
-                services=services,
-                deep_sleep=(addr[0] in self._unicasts),
-                model=_get_model(services),
+                services=query_resp.parser.parse(),
+                deep_sleep=query_resp.deep_sleep,
+                model=_get_model(query_resp.parser.parse()),
             )
 
             if self.end_condition(response):
                 # Matches end condition: replace everything found so far and abort
-                self.responses = {IPv4Address(addr[0]): response}
+                self.query_responses = {addr[0]: self.query_responses[addr[0]]}
                 self.semaphore.release()
                 self.close()
-            else:
-                self.responses[IPv4Address(addr[0])] = response
 
     @staticmethod
     def error_received(exc) -> None:
@@ -436,7 +507,7 @@ async def multicast(  # pylint: disable=too-many-arguments
     port: int = 5353,
     timeout: int = 4,
     end_condition: typing.Optional[typing.Callable[[Response], bool]] = None,
-) -> typing.Dict[IPv4Address, Response]:
+) -> typing.List[Response]:
     """Send multicast request for services."""
     protocol = MulticastDnsSdClientProtocol(
         loop, services, address, port, end_condition
diff --git a/pyatv/support/scan.py b/pyatv/support/scan.py
index 981638dd..b4599a1f 100644
--- a/pyatv/support/scan.py
+++ b/pyatv/support/scan.py
@@ -20,6 +20,7 @@ ScanHandler = Callable[[mdns.Service, mdns.Response], Optional[ScanHandlerReturn
 ScanMethod = Callable[[], Mapping[str, ScanHandler]]
 
 DEVICE_INFO: str = "_device-info._tcp.local"
+SLEEP_PROXY: str = "_sleep-proxy._udp.local"
 
 # These ports have been "arbitrarily" chosen (see issue #580) because a device normally
 # listen on them (more or less). They are used as best-effort when for unicast scanning
@@ -54,11 +55,16 @@ class BaseScanner(ABC):
     def __init__(self) -> None:
         """Initialize a new BaseScanner."""
         self._services: Dict[str, ScanHandler] = {
-            DEVICE_INFO: lambda service, response: None
+            DEVICE_INFO: self._empty_handler,
+            SLEEP_PROXY: self._empty_handler,
         }
         self._found_devices: Dict[IPv4Address, FoundDevice] = {}
         self._properties: Dict[IPv4Address, Dict[str, Mapping[str, str]]] = {}
 
+    @staticmethod
+    def _empty_handler(service: mdns.Service, response: mdns.Response) -> None:
+        pass
+
     def add_service(self, service_type: str, handler: ScanHandler):
         """Add service type to discover."""
         self._services[service_type] = handler
@@ -200,7 +206,7 @@ class MulticastMdnsScanner(BaseScanner):
             timeout=timeout,
             end_condition=self._end_if_identifier_found if self.identifier else None,
         )
-        for _, response in responses.items():
+        for response in responses:
             self.handle_response(response)
 
     def _end_if_identifier_found(self, response: mdns.Response):


Test patch:
diff --git a/tests/fake_udns.py b/tests/fake_udns.py
index d024a6f7..ae2941a3 100644
--- a/tests/fake_udns.py
+++ b/tests/fake_udns.py
@@ -282,7 +282,7 @@ def stub_multicast(udns_server, loop):
         hosts = set(
             chain(*[service.addresses for service in udns_server.services.values()])
         )
-        devices = {}
+        devices = []
         sleep_proxy = udns_server.sleep_proxy
         udns_server.sleep_proxy = False
         for host in hosts:
@@ -290,8 +290,8 @@ def stub_multicast(udns_server, loop):
             response = await mdns.unicast(
                 loop, "127.0.0.1", services, port=udns_server.port
             )
-            devices[IPv4Address(host)] = mdns.Response(
-                response.services, sleep_proxy, response.model
+            devices.append(
+                mdns.Response(response.services, sleep_proxy, response.model)
             )
         return devices
 
diff --git a/tests/support/test_mdns.py b/tests/support/test_mdns.py
index def4b872..10722166 100644
--- a/tests/support/test_mdns.py
+++ b/tests/support/test_mdns.py
@@ -4,7 +4,7 @@ import asyncio
 import io
 from ipaddress import IPv4Address
 import logging
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import pytest
 
@@ -30,21 +30,27 @@ TEST_SERVICES = dict(
 def get_response_for_service(
     service: str,
 ) -> Tuple[dns.DnsMessage, Optional[fake_udns.FakeDnsService]]:
-    req = mdns.create_request([service])
+    req = mdns.create_service_queries([service], mdns.QueryType.PTR)[0]
     resp = fake_udns.create_response(req, TEST_SERVICES)
     return dns.DnsMessage().unpack(resp.pack()), TEST_SERVICES.get(service)
 
 
+def parse_services(message: mdns.DnsMessage) -> List[mdns.Service]:
+    parser = mdns.ServiceParser()
+    parser.add_message(message)
+    return parser.parse()
+
+
 def test_non_existing_service():
     resp, _ = get_response_for_service("_missing")
-    assert len(resp.questions) == 1
+    assert len(resp.questions) == 2
     assert len(resp.answers) == 0
     assert len(resp.resources) == 0
 
 
 def test_service_has_expected_responses():
     resp, _ = get_response_for_service(MEDIAREMOTE_SERVICE)
-    assert len(resp.questions) == 1
+    assert len(resp.questions) == 2
     assert len(resp.answers) == 1
     assert len(resp.resources) == 3
 
@@ -127,14 +133,14 @@ def test_authority():
 
 
 def test_parse_empty_service():
-    assert mdns.parse_services(dns.DnsMessage()) == []
+    assert parse_services(dns.DnsMessage()) == []
 
 
 def test_parse_no_service_type():
     service_params = (None, "service", [], 0, {})
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 0
 
 
@@ -142,7 +148,7 @@ def test_parse_no_service_name():
     service_params = ("_abc._tcp.local", None, [], 0, {})
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 0
 
 
@@ -150,7 +156,7 @@ def test_parse_with_name_and_type():
     service_params = ("_abc._tcp.local", "service", [], 0, {})
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 1
     dns_utils.assert_service(parsed[0], *service_params)
 
@@ -159,7 +165,7 @@ def test_parse_with_port_and_address():
     service_params = ("_abc._tcp.local", "service", ["10.0.0.1"], 123, {})
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 1
     dns_utils.assert_service(parsed[0], *service_params)
 
@@ -168,7 +174,7 @@ def test_parse_single_service():
     service_params = ("_abc._tcp.local", "service", ["10.0.10.1"], 123, {"foo": "bar"})
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 1
     dns_utils.assert_service(parsed[0], *service_params)
 
@@ -191,7 +197,7 @@ def test_parse_double_service():
     message = dns_utils.add_service(dns.DnsMessage(), *service1_params)
     message = dns_utils.add_service(message, *service2_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 2
     dns_utils.assert_service(parsed[0], *service1_params)
     dns_utils.assert_service(parsed[1], *service2_params)
@@ -208,7 +214,7 @@ def test_parse_pick_one_available_address():
     )
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 1
     assert str(parsed[0].address) in addresses
 
@@ -223,7 +229,7 @@ def test_parse_ignore_link_local_address():
     )
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 1
     assert parsed[0].address is None
 
@@ -236,7 +242,24 @@ def test_parse_properties_converts_keys_to_lower_case():
     service_params = ("_abc._tcp.local", "service", [], 0, {"FOO": "bar", "Bar": "FOO"})
     message = dns_utils.add_service(dns.DnsMessage(), *service_params)
 
-    parsed = mdns.parse_services(message)
+    parsed = parse_services(message)
     assert len(parsed) == 1
     assert parsed[0].properties["foo"] == "bar"
     assert parsed[0].properties["Bar"] == "FOO"
+
+
+def test_parse_ignore_duplicate_records():
+    service_params = ("_abc._tcp.local", "service", [], 0, {})
+    message = dns_utils.add_service(dns.DnsMessage(), *service_params)
+
+    parser = mdns.ServiceParser()
+    parser.add_message(message)
+    parser.add_message(message)
+
+    # One service should be present in the table
+    assert len(parser.table) == 1
+
+    # A single record should be there since duplicates is ignored
+    records = parser.table["service._abc._tcp.local"]
+    assert mdns.QueryType.SRV in records
+    assert len(records[mdns.QueryType.SRV]) == 1
diff --git a/tests/support/test_mdns_functional.py b/tests/support/test_mdns_functional.py
index 639144dd..5baa6f69 100644
--- a/tests/support/test_mdns_functional.py
+++ b/tests/support/test_mdns_functional.py
@@ -16,6 +16,7 @@ SERVICE_NAME = "Kitchen"
 MEDIAREMOTE_SERVICE = "_mediaremotetv._tcp.local"
 DEVICE_INFO_SERVICE = "_device-info._tcp._local"
 
+SERVICES_PER_REQUEST = 3
 
 TEST_SERVICES = dict(
     [
@@ -26,6 +27,13 @@ TEST_SERVICES = dict(
 )
 
 
+def gen_test_services(count: int) -> typing.List[str]:
+    return [f"srv{i}._tcp.local" for i in range(count)]
+
+
+pytestmark = pytest.mark.asyncio
+
+
 @pytest.fixture(autouse=True)
 def mdns_debug():
     logger = logging.getLogger("pyatv.support.mdns")
@@ -75,7 +83,7 @@ async def multicast_fastexit(event_loop, monkeypatch, udns_server):
 
     # Checks if either response or request count has been fulfilled
     def _check_cond(protocol: mdns.MulticastDnsSdClientProtocol) -> bool:
-        if len(protocol.responses) == conditions["responses"]:
+        if len(protocol.query_responses) == conditions["responses"]:
             return False
         if conditions["requests"] == 0:
             return True
@@ -97,42 +105,55 @@ async def multicast_fastexit(event_loop, monkeypatch, udns_server):
     await asyncio.gather(*clients)
 
 
-async def unicast(event_loop, udns_server, service_name, timeout=1):
+async def unicast(event_loop, udns_server, service_names, timeout=1):
     return (
         await mdns.unicast(
             event_loop,
             "127.0.0.1",
-            [service_name],
+            service_names,
             port=udns_server.port,
             timeout=timeout,
         ),
-        TEST_SERVICES.get(service_name),
+        TEST_SERVICES.get(service_names[0]),
     )
 
 
-@pytest.mark.asyncio
 async def test_unicast_has_valid_service(event_loop, udns_server):
-    resp, service = await unicast(event_loop, udns_server, MEDIAREMOTE_SERVICE)
+    resp, service = await unicast(event_loop, udns_server, [MEDIAREMOTE_SERVICE])
     assert len(resp.services) == 1
     assert resp.services[0].type == MEDIAREMOTE_SERVICE
     assert resp.services[0].name == service.name
     assert resp.services[0].port == service.port
 
 
-@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "service_count,expected_requests",
+    [
+        (1, 1),
+        (SERVICES_PER_REQUEST, 1),
+        (SERVICES_PER_REQUEST + 1, 2),
+        (2 * SERVICES_PER_REQUEST + 1, 3),
+    ],
+)
+async def test_unicast_multiple_requests(
+    service_count, expected_requests, event_loop, udns_server
+):
+    resp, _ = await unicast(event_loop, udns_server, gen_test_services(service_count))
+    assert udns_server.request_count == expected_requests
+
+
 async def test_unicast_resend_if_no_response(event_loop, udns_server):
     udns_server.skip_count = 2
-    resp, service = await unicast(event_loop, udns_server, MEDIAREMOTE_SERVICE, 3)
+    resp, service = await unicast(event_loop, udns_server, [MEDIAREMOTE_SERVICE], 3)
     assert len(resp.services) == 1
     assert resp.services[0].type == MEDIAREMOTE_SERVICE
     assert resp.services[0].name == service.name
     assert resp.services[0].port == service.port
 
 
-@pytest.mark.asyncio
 async def test_unicast_specific_service(event_loop, udns_server):
     resp, _ = await unicast(
-        event_loop, udns_server, SERVICE_NAME + "." + MEDIAREMOTE_SERVICE
+        event_loop, udns_server, [SERVICE_NAME + "." + MEDIAREMOTE_SERVICE]
     )
     assert len(resp.services) == 1
 
@@ -142,14 +163,36 @@ async def test_unicast_specific_service(event_loop, udns_server):
     assert resp.services[0].port == service.port
 
 
-@pytest.mark.asyncio
+async def test_unicast_includes_sleep_proxy_service(event_loop, udns_server):
+    udns_server.services = {
+        "_test._tcp.local": fake_udns.FakeDnsService(
+            name="test", addresses=["127.0.0.1"], port=1234, properties={}, model=None
+        ),
+        "_sleep-proxy._udp.local": fake_udns.FakeDnsService(
+            name="sleepy", addresses=["127.0.0.1"], port=5678, properties={}, model=None
+        ),
+    }
+
+    # _sleep-proxy._udp.local should be requested implicitly in unicast for
+    # any service
+    resp, _ = await unicast(event_loop, udns_server, ["_test._tcp.local"])
+    assert len(resp.services) == 2
+
+    proxy = [
+        service
+        for service in resp.services
+        if service.type == "_sleep-proxy._udp.local"
+    ][0]
+    assert proxy.name == "sleepy"
+    assert proxy.port == 5678
+
+
 async def test_multicast_no_response(event_loop, udns_server, multicast_fastexit):
     multicast_fastexit(responses=0, requests=0)
 
     await mdns.multicast(event_loop, [], "127.0.0.1", udns_server.port)
 
 
-@pytest.mark.asyncio
 async def test_multicast_has_valid_service(event_loop, udns_server, multicast_fastexit):
     multicast_fastexit(responses=1, requests=0)
 
@@ -158,13 +201,12 @@ async def test_multicast_has_valid_service(event_loop, udns_server, multicast_fa
     )
     assert len(resp) == 1
 
-    first = resp[IPv4Address("127.0.0.1")].services[0]
+    first = resp[0].services[0]
     assert first.type == MEDIAREMOTE_SERVICE
     assert first.name == SERVICE_NAME
     assert first.port == 1234
 
 
-@pytest.mark.asyncio
 async def test_multicast_end_condition_met(
     event_loop, udns_server, multicast_fastexit, stub_ip_address
 ):
@@ -184,10 +226,9 @@ async def test_multicast_end_condition_met(
         end_condition=_end_cond,
     )
     assert len(resp) == 1
-    actor.assert_called_once_with(resp[IPv4Address("127.0.0.1")])
+    actor.assert_called_once_with(resp[0])
 
 
-@pytest.mark.asyncio
 async def test_multicast_sleeping_device(event_loop, udns_server, multicast_fastexit):
     multicast_fastexit(responses=0, requests=3)
     udns_server.sleep_proxy = True
@@ -212,7 +253,6 @@ async def test_multicast_sleeping_device(event_loop, udns_server, multicast_fast
     assert len(resp) == 1
 
 
-@pytest.mark.asyncio
 async def test_multicast_deep_sleep(event_loop, udns_server, multicast_fastexit):
     multicast_fastexit(responses=1, requests=0)
 
@@ -220,7 +260,7 @@ async def test_multicast_deep_sleep(event_loop, udns_server, multicast_fastexit)
         event_loop, [MEDIAREMOTE_SERVICE], "127.0.0.1", udns_server.port
     )
 
-    assert not resp[IPv4Address("127.0.0.1")].deep_sleep
+    assert not resp[0].deep_sleep
 
     udns_server.sleep_proxy = True
     multicast_fastexit(responses=1, requests=0)
@@ -229,10 +269,9 @@ async def test_multicast_deep_sleep(event_loop, udns_server, multicast_fastexit)
         event_loop, [MEDIAREMOTE_SERVICE], "127.0.0.1", udns_server.port
     )
 
-    assert resp[IPv4Address("127.0.0.1")].deep_sleep
+    assert resp[0].deep_sleep
 
 
-@pytest.mark.asyncio
 async def test_multicast_device_model(event_loop, udns_server, multicast_fastexit):
     multicast_fastexit(responses=1, requests=0)
 
@@ -240,7 +279,7 @@ async def test_multicast_device_model(event_loop, udns_server, multicast_fastexi
         event_loop, [MEDIAREMOTE_SERVICE], "127.0.0.1", udns_server.port
     )
 
-    assert not resp[IPv4Address("127.0.0.1")].model
+    assert not resp[0].model
 
     udns_server.services = {
         MEDIAREMOTE_SERVICE: fake_udns.FakeDnsService(
@@ -257,4 +296,4 @@ async def test_multicast_device_model(event_loop, udns_server, multicast_fastexi
         event_loop, [MEDIAREMOTE_SERVICE], "127.0.0.1", udns_server.port
     )
 
-    assert resp[IPv4Address("127.0.0.1")].model == "dummy"
+    assert resp[0].model == "dummy"
diff --git a/tests/support/test_scan.py b/tests/support/test_scan.py
index 45cf4b14..51db16a6 100644
--- a/tests/support/test_scan.py
+++ b/tests/support/test_scan.py
@@ -10,8 +10,6 @@ from pyatv.support.scan import BaseScanner, get_unique_identifiers
 TEST_SERVICE1 = Service("_service1._tcp.local", "service1", None, 0, {"a": "b"})
 TEST_SERVICE2 = Service("_service2._tcp.local", "service2", None, 0, {"c": "d"})
 
-# class TestScanner(BaseScanner):
-
 
 @pytest.fixture
 def response():