Raise better error if server doesn't return a full answer

mcstatus/protocol/connection.py

@@ -460,7 +460,10 @@ def __init__(self) -> None:
         self.received = bytearray()
 
     def read(self, length: int) -> bytearray:
-        """Return self.received up to length bytes, then cut received up to that point."""
+        """Return :attr:`.received` up to length bytes, then cut received up to that point."""
+        if len(self.received) < length:
+            raise IOError(f"Not enough data to read! {len(self.received)} < {length}")
+
         result = self.received[:length]
         self.received = self.received[length:]
         return result

tests/protocol/test_connection.py

@@ -204,6 +204,18 @@ def test_write_buffer(self):
 
         assert self.connection.flush() == bytearray.fromhex("027FAA")
 
+    def test_read_empty(self):
+        self.connection.received = bytearray()
+
+        with pytest.raises(IOError):
+            self.connection.read(1)
+
+    def test_read_not_enough(self):
+        self.connection.received = bytearray(b"a")
+
+        with pytest.raises(IOError):
+            self.connection.read(2)
+
 
 class TestTCPSocketConnection:
     @pytest.fixture(scope="class")
@@ -236,7 +248,13 @@ def test_read(self, connection):
         assert connection.read(2) == bytearray.fromhex("7FAA")
 
     def test_read_empty(self, connection):
-        connection.socket.recv.return_value = bytearray.fromhex("")
+        connection.socket.recv.return_value = bytearray()
+
+        with pytest.raises(IOError):
+            connection.read(1)
+
+    def test_read_not_enough(self, connection):
+        connection.socket.recv.side_effect = [bytearray(b"a"), bytearray()]
 
         with pytest.raises(IOError):
             connection.read(2)
@@ -276,6 +294,12 @@ def test_read(self, connection):
 
         assert connection.read(2) == bytearray.fromhex("7FAA")
 
+    def test_read_empty(self, connection):
+        connection.socket.recvfrom.return_value = []
+
+        with pytest.raises(IndexError):
+            connection.read(1)
+
     def test_write(self, connection):
         connection.write(bytearray.fromhex("7FAA"))