From 8e9c1da39c16556fdab75f8a5dd6a031ec1bcdc0 Mon Sep 17 00:00:00 2001
From: Floris272 <florispuijk@outlook.com>
Date: Fri, 22 Mar 2024 15:10:08 +0100
Subject: [PATCH] Add on_open parameter to Client.connect

---
 pyproject.toml                | 2 +-
 src/bluecurrent_api/client.py | 8 ++++----
 tests/test_client.py          | 6 +++++-
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index a6840e0..48edcd1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bluecurrent-api"
-version = "1.2.2"
+version = "1.2.3"
 authors = [
   { name="Floris272", email="florispuijk@outlook.com" },
 ]
diff --git a/src/bluecurrent_api/client.py b/src/bluecurrent_api/client.py
index 585cb3b..c97f72b 100644
--- a/src/bluecurrent_api/client.py
+++ b/src/bluecurrent_api/client.py
@@ -34,7 +34,7 @@ async def get_email(self) -> str:
         """Get user email."""
         return await self.websocket.get_email()
 
-    async def _on_open(self) -> None:
+    async def _on_open(self, on_open: Callable[[], Coroutine[Any, Any, None]],) -> None:
         """Send requests when connected."""
         await self.websocket.send_request(
             {
@@ -42,8 +42,7 @@ async def _on_open(self) -> None:
                 "header": "homeassistant",
             }
         )
-        await self.get_charge_cards()
-        await self.get_charge_points()
+        await on_open()
 
     def get_next_reset_delta(self) -> timedelta:
         """Returns the timedelta until the websocket limits are reset."""
@@ -52,9 +51,10 @@ def get_next_reset_delta(self) -> timedelta:
     async def connect(
         self,
         receiver: Callable[[dict[str, Any]], Coroutine[Any, Any, None]],
+        on_open: Callable[[], Coroutine[Any, Any, None]],
     ) -> None:
         """Connect to the websocket."""
-        await self.websocket.start(receiver, self._on_open)
+        await self.websocket.start(receiver, lambda : self._on_open(on_open))
 
     async def disconnect(self) -> None:
         """Disconnect the websocket."""
diff --git a/tests/test_client.py b/tests/test_client.py
index b6f7abe..2c2e77b 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -98,7 +98,11 @@ async def test_on_open(mocker: MockerFixture):
     )
     client = Client()
 
-    await client._on_open()
+    async def on_open():
+        await client.get_charge_cards()
+        await client.get_charge_points()
+
+    await client._on_open(on_open)
     test_send_request.assert_has_calls(
         [
             mocker.call(