diff --git a/serfclient/client.py b/serfclient/client.py index a98d10d..97ffdfe 100644 --- a/serfclient/client.py +++ b/serfclient/client.py @@ -74,3 +74,11 @@ def stats(self): Obtain operator debugging information about the running Serf agent. """ return self.connection.call('stats') + + def close(self): + """ + Close connection to Serf agent. + """ + if self.connection: + self.connection.close() + self.connection = None diff --git a/serfclient/connection.py b/serfclient/connection.py index 5529037..4a38b86 100644 --- a/serfclient/connection.py +++ b/serfclient/connection.py @@ -166,3 +166,11 @@ def _decode_addr_key(self, obj_dict): obj_dict[key] = ip_addr.encode('utf-8') return obj_dict + + def close(self): + """ + Close the connection with the Serf agent. + """ + if self._socket: + self._socket.close() + self._socket = None diff --git a/tests/test_client.py b/tests/test_client.py index f85ca45..ee42a53 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -110,3 +110,8 @@ def test_stats_is_well_formed(self, serf): for key in [b'agent', b'runtime', b'serf', b'tags']: assert key in stats.body assert isinstance(stats.body[key], dict) + + def test_close(self, serf): + assert serf.connection is not None + serf.close() + assert serf.connection is None diff --git a/tests/test_connection.py b/tests/test_connection.py index b4f49cd..45cd7d6 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -134,3 +134,9 @@ def test_decode_addr_key_ipv4_mapped_ipv6(self, rpc): def test_decode_addr_key_ipv4(self, rpc): ip_address = '192.168.0.1' assert extract_addr(rpc, ip_address, socket.AF_INET) == ip_address + + def test_close(self, rpc): + rpc.handshake() + assert rpc._socket is not None + rpc.close() + assert rpc._socket is None