From 08f23715d963ccdb712133fdceb9e35f62b76561 Mon Sep 17 00:00:00 2001
From: Wonko <wonko@hanstool.org>
Date: Fri, 17 Dec 2021 12:35:01 +0100
Subject: [PATCH] MPDClient: add __enter__/__exit__ to connect/disconnect

---
 musicpd.py |  7 +++++++
 test.py    | 12 ++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/musicpd.py b/musicpd.py
index d96e76a..5824056 100644
--- a/musicpd.py
+++ b/musicpd.py
@@ -726,6 +726,13 @@ class MPDClient:
             self._sock.close()
         self._reset()
 
+    def __enter__(self):
+        self.connect()
+        return self
+
+    def __exit__(self, exception_type, exception_value, exception_traceback):
+        self.disconnect()
+
     def fileno(self):
         """Return the socket’s file descriptor (a small integer).
         This is useful with :py:obj:`select.select`.
diff --git a/test.py b/test.py
index 2c5094d..c911a57 100755
--- a/test.py
+++ b/test.py
@@ -593,6 +593,18 @@ class testException(unittest.TestCase):
             with self.assertRaises(musicpd.CommandError):
                 cli.find('(album == "foo\nbar")')
 
+class testContextManager(unittest.TestCase):
+
+    def test_enter_exit(self):
+        os.environ['MPD_HOST'] = '@abstract'
+        with mock.patch('musicpd.socket') as socket_mock:
+            sock = mock.MagicMock(name='socket')
+            socket_mock.socket.return_value = sock
+            cli = musicpd.MPDClient()
+            with cli as c:
+                sock.connect.assert_called_with('\0abstract')
+                sock.close.assert_not_called()
+            sock.close.assert_called()
 
 if __name__ == '__main__':
     unittest.main()
-- 
2.39.2