From 8e42f911092f647438c5bcab711b8bd1619ab3ed Mon Sep 17 00:00:00 2001
From: Renaud Bruyeron <bruyeron@gmail.com>
Date: Sun, 21 Jan 2024 18:19:14 +0100
Subject: [PATCH] Add a simpler entry point that bypasses the discovery process
 in the case where the Inverter type is already known.

Some inverters at the end of the REGISTRY (e.g. X1HybridGen4) randomly
crash because the discovery code issues too many incorrect requests
before sending the correct one.
This patch adds a variant of the entry point in which one can specify
the name of the inverter type.

This addresses https://github.com/home-assistant/core/issues/66617
and https://github.com/home-assistant/core/issues/99421
---
 solax/__init__.py       |  8 +++++++-
 solax/discovery.py      | 37 +++++++++++++++++++++++++++++++++++++
 tests/test_discovery.py | 19 +++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/solax/__init__.py b/solax/__init__.py
index 5953b06..a13385c 100644
--- a/solax/__init__.py
+++ b/solax/__init__.py
@@ -4,7 +4,7 @@
 
 from async_timeout import timeout
 
-from solax.discovery import discover
+from solax.discovery import discover, discover_single_inverter
 from solax.inverter import Inverter, InverterResponse
 
 _LOGGER = logging.getLogger(__name__)
@@ -36,6 +36,12 @@ async def real_time_api(ip_address, port=80, pwd=""):
     return RealTimeAPI(i)
 
 
+# bypass the full discovery process, and only do it for the specified inverter type
+async def real_time_api_inverter(inverter_name, ip_address, port=80, pwd=""):
+    i = await discover_single_inverter(inverter_name, ip_address, port, pwd)
+    return RealTimeAPI(i)
+
+
 class RealTimeAPI:
     """Solax inverter real time API"""
 
diff --git a/solax/discovery.py b/solax/discovery.py
index 9cfdd60..9d62d01 100644
--- a/solax/discovery.py
+++ b/solax/discovery.py
@@ -92,6 +92,37 @@ async def discover(self, host, port, pwd="") -> Inverter:
         )
         raise DiscoveryError(msg)
 
+    async def discover_single_inverter(
+        self, inverter_name, host, port, pwd=""
+    ) -> Inverter:
+        found = False
+        for invclass in REGISTRY:
+            if invclass.__name__ == inverter_name:
+                found = True
+                for i in invclass.build_all_variants(host, port, pwd):
+                    task = asyncio.create_task(self._discovery_task(i), name=f"{i}")
+                    task.add_done_callback(self._task_handler)
+                    self._tasks.add(task)
+
+        if not found:
+            raise DiscoveryError(f"Unknown inverter type {inverter_name}")
+
+        while len(self._tasks) > 0:
+            logging.debug("%d discovery tasks are still running...", len(self._tasks))
+            await asyncio.sleep(0.5)
+
+        if self._discovered_inverter is not None:
+            logging.info("Discovered inverter: %s", self._discovered_inverter)
+            return self._discovered_inverter
+
+        msg = (
+            "Unable to connect to the inverter at "
+            f"host={host} port={port} of type {inverter_name}.\n"
+            "Please see https://github.com/squishykid/solax/wiki/DiscoveryError\n"
+            f"Failures={str(self._failures)}"
+        )
+        raise DiscoveryError(msg)
+
 
 class DiscoveryError(Exception):
     """Raised when unable to discover inverter"""
@@ -101,3 +132,9 @@ async def discover(host, port, pwd="") -> Inverter:
     discover_state = DiscoveryState()
     await discover_state.discover(host, port, pwd)
     return discover_state.get_discovered_inverter()
+
+
+async def discover_single_inverter(inverter, host, port, pwd="") -> Inverter:
+    discover_state = DiscoveryState()
+    await discover_state.discover_single_inverter(inverter, host, port, pwd)
+    return discover_state.get_discovered_inverter()
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index 2f57b50..5246c30 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -11,12 +11,31 @@ async def test_discovery(inverters_fixture):
     assert rt_api.inverter.__class__ == inverter_class
 
 
+@pytest.mark.asyncio
+async def test_discovery_single_inverter(inverters_fixture):
+    conn, inverter_class, _ = inverters_fixture
+    rt_api = await solax.real_time_api_inverter(inverter_class.__name__, *conn)
+    assert rt_api.inverter.__class__ == inverter_class
+
+
+@pytest.mark.asyncio
+async def test_discovery_unsupported_inverter():
+    with pytest.raises(DiscoveryError):
+        await solax.real_time_api_inverter("doesnotexist", "localhost", 2)
+
+
 @pytest.mark.asyncio
 async def test_discovery_no_host():
     with pytest.raises(DiscoveryError):
         await solax.real_time_api("localhost", 2)
 
 
+@pytest.mark.asyncio
+async def test_discovery_single_inverter_no_host():
+    with pytest.raises(DiscoveryError):
+        await solax.real_time_api_inverter("X1HybridGen4", "localhost", 2)
+
+
 @pytest.mark.asyncio
 async def test_discovery_no_host_with_pwd():
     with pytest.raises(DiscoveryError):