diff --git a/realtime/connection.py b/realtime/connection.py index 5ee9e8b..b87c014 100644 --- a/realtime/connection.py +++ b/realtime/connection.py @@ -113,6 +113,18 @@ async def _connect(self) -> None: else: raise Exception("Connection Failed") + @ensure_connection + def close(self) -> None: + """ + Wrapper for async def _close() to expose a non-async interface + """ + loop = asyncio.get_event_loop() + loop.run_until_complete(self._close()) + self.connected = False + + async def _close(self) -> None: + await self.ws_connection.close() + async def _keep_alive(self) -> None: """ Sending heartbeat to server every 5 seconds diff --git a/tests/tests_close.py b/tests/tests_close.py new file mode 100644 index 0000000..e6f6002 --- /dev/null +++ b/tests/tests_close.py @@ -0,0 +1,27 @@ +import os + +from realtime.connection import Socket + + +def callback(payload): + print("Callback:\t", payload) + + +def test(): + SUPABASE_ID = os.environ.get("SUPABASE_TEST_URL") + API_KEY = os.environ.get("SUPABASE_TEST_KEY") + + URL = f"wss://{SUPABASE_ID}.supabase.co/realtime/v1/websocket?apikey={API_KEY}&vsn=1.0.0" + + s = Socket(URL) + s.connect() + + channel_1 = s.set_channel("realtime:public:sample") + channel_1.join().on("INSERT", callback) + + s.close() + assert s.ws_connection.closed, "Connection was not closed as expected." + + +if __name__ == "__main__": + test()