[dns] minor fixes

mitmproxy/net/udp.py

@@ -217,7 +217,13 @@ class DatagramReader:
 
     async def read(self, n: int) -> bytes:
         assert n >= MAX_DATAGRAM_SIZE
-        return b'' if self._eof else await self._packets.get()
+        if self._eof:
+            try:
+                return self._packets.get_nowait()
+            except asyncio.QueueEmpty:
+                return b''
+        else:
+            return await self._packets.get()
 
 
 class DatagramWriter:

mitmproxy/test/tflow.py

@@ -114,6 +114,7 @@ def tdnsflow(
     assert err is False or isinstance(err, flow.Error)
 
     f = dns.DNSFlow(client_conn, server_conn)
+    f.timestamp_created = req.timestamp
     f.request = req
     f.response = resp or None
     f.error = err or None

test/mitmproxy/net/test_udp.py

@@ -8,12 +8,12 @@ async def test_reader():
     reader = DatagramReader()
     addr = ('8.8.8.8', 53)
     reader.feed_data(b'First message', addr)
-    with pytest.raises():
-        reader.feed_data(bytearray(MAX_DATAGRAM_SIZE + 1))
+    with pytest.raises(AssertionError):
+        reader.feed_data(bytearray(MAX_DATAGRAM_SIZE + 1), addr)
     reader.feed_data(b'Second message', addr)
     reader.feed_eof()
     assert await reader.read(65535) == b'First message'
-    with pytest.raises():
+    with pytest.raises(AssertionError):
         await reader.read(MAX_DATAGRAM_SIZE - 1)
     assert await reader.read(65535) == b'Second message'
     assert not await reader.read(65535)