From 8e046bb25dee895fb3f370f55209c3d714497288 Mon Sep 17 00:00:00 2001
From: Luis Sousa <luis.m.sousa@outlook.pt>
Date: Thu, 8 Feb 2024 16:57:01 +0000
Subject: [PATCH] re-submitting @rowinho 's PR to trigger the tests

---
 solax/__init__.py  |  4 ++--
 solax/discovery.py | 20 +++++++++++++-------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/solax/__init__.py b/solax/__init__.py
index 6fd9817..d259b98 100644
--- a/solax/__init__.py
+++ b/solax/__init__.py
@@ -32,8 +32,8 @@ async def rt_request(inv: Inverter, retry, t_wait=0) -> InverterResponse:
         raise
 
 
-async def real_time_api(ip_address, port=80, pwd=""):
-    i = await discover(ip_address, port, pwd)
+async def real_time_api(ip_address, port=80, pwd="", model=""):
+    i = await discover(ip_address, port, pwd, model)
     return RealTimeAPI(i)
 
 
diff --git a/solax/discovery.py b/solax/discovery.py
index 9cfdd60..abaf3e3 100644
--- a/solax/discovery.py
+++ b/solax/discovery.py
@@ -69,12 +69,13 @@ async def _discovery_task(cls, i) -> Inverter:
         await i.get_data()
         return i
 
-    async def discover(self, host, port, pwd="") -> Inverter:
+    async def discover(self, host, port, pwd="", model=None) -> Inverter:
         for inverter in REGISTRY:
-            for i in inverter.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 inverter.__name__ == model or model is None:
+                for i in inverter.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)
 
         while len(self._tasks) > 0:
             logging.debug("%d discovery tasks are still running...", len(self._tasks))
@@ -97,7 +98,12 @@ class DiscoveryError(Exception):
     """Raised when unable to discover inverter"""
 
 
-async def discover(host, port, pwd="") -> Inverter:
+async def discover(host, port, pwd="", model=None) -> Inverter:
     discover_state = DiscoveryState()
-    await discover_state.discover(host, port, pwd)
+    await discover_state.discover(host, port, pwd, model)
     return discover_state.get_discovered_inverter()
+
+def get_models() -> List[str]:
+    models = list(map(lambda inverter: inverter.__name__, REGISTRY))
+    models.sort()
+    return models