diff --git a/mcstatus/server.py b/mcstatus/server.py index c8eab3a..2b12f2e 100644 --- a/mcstatus/server.py +++ b/mcstatus/server.py @@ -1,5 +1,3 @@ -import json - from mcstatus.pinger import ServerPinger from mcstatus.protocol.connection import TCPSocketConnection, UDPSocketConnection from mcstatus.querier import ServerQuerier @@ -10,14 +8,30 @@ class MinecraftServer: self.host = host self.port = port - def ping_server(self, **kwargs): + def ping_server(self, retries=3, **kwargs): + attempt = 0 connection = TCPSocketConnection((self.host, self.port)) - pinger = ServerPinger(connection, host=self.host, port=self.port, **kwargs) - pinger.handshake() - return pinger.read_status(), pinger.test_ping() + exception = None + while attempt < retries: + try: + pinger = ServerPinger(connection, host=self.host, port=self.port, **kwargs) + pinger.handshake() + return pinger.read_status(), pinger.test_ping() + except Exception as e: + exception = e + attempt += 1 + raise exception - def query_server(self): - connection = UDPSocketConnection((self.host, self.port)) - querier = ServerQuerier(connection) - querier.handshake() - return querier.read_query() \ No newline at end of file + def query_server(self, retries=3): + attempt = 0 + exception = None + while attempt < retries: + try: + connection = UDPSocketConnection((self.host, self.port)) + querier = ServerQuerier(connection) + querier.handshake() + return querier.read_query() + except Exception as e: + exception = e + attempt += 1 + raise exception \ No newline at end of file diff --git a/mcstatus/tests/test_server.py b/mcstatus/tests/test_server.py index 56bf6d5..b7ce100 100644 --- a/mcstatus/tests/test_server.py +++ b/mcstatus/tests/test_server.py @@ -23,6 +23,14 @@ class TestMinecraftServer(TestCase): self.assertEqual(info.raw, {"description":"A Minecraft Server","players":{"max":20,"online":0},"version":{"name":"1.8","protocol":47}}) self.assertTrue(latency >= 0) + def test_ping_server_retry(self): + with patch("mcstatus.server.TCPSocketConnection") as connection: + connection.return_value = None + with patch("mcstatus.server.ServerPinger") as pinger: + pinger.side_effect = [Exception, Exception, Exception] + self.assertRaises(Exception, self.server.ping_server) + self.assertEqual(pinger.call_count, 3) + def test_query_server(self): self.socket.receive(bytearray.fromhex("090000000035373033353037373800")) self.socket.receive(bytearray.fromhex("00000000000000000000000000000000686f73746e616d650041204d696e656372616674205365727665720067616d657479706500534d500067616d655f6964004d494e4543524146540076657273696f6e00312e3800706c7567696e7300006d617000776f726c64006e756d706c61796572730033006d6178706c617965727300323000686f7374706f727400323535363500686f73746970003139322e3136382e35362e31000001706c617965725f000044696e6e6572626f6e6500446a696e6e69626f6e650053746576650000")) @@ -46,4 +54,12 @@ class TestMinecraftServer(TestCase): "maxplayers": "20", "hostport": "25565", "hostip": "192.168.56.1", - }) \ No newline at end of file + }) + + def test_query_server_retry(self): + with patch("mcstatus.server.UDPSocketConnection") as connection: + connection.return_value = None + with patch("mcstatus.server.ServerQuerier") as querier: + querier.side_effect = [Exception, Exception, Exception] + self.assertRaises(Exception, self.server.query_server) + self.assertEqual(querier.call_count, 3) \ No newline at end of file