diff --git a/solax/__init__.py b/solax/__init__.py index 6fd9817..4a721ac 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=None): + i = await discover(ip_address, port, pwd, model) return RealTimeAPI(i) diff --git a/solax/discovery.py b/solax/discovery.py index 9cfdd60..7599569 100644 --- a/solax/discovery.py +++ b/solax/discovery.py @@ -1,6 +1,7 @@ import asyncio import logging import typing +from typing import List from solax.inverter import Inverter, InverterError from solax.inverters import ( @@ -69,12 +70,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 +99,13 @@ 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 diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 2f57b50..26c2052 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -11,6 +11,20 @@ async def test_discovery(inverters_fixture): assert rt_api.inverter.__class__ == inverter_class +@pytest.mark.asyncio +async def test_discovery_with_model(inverters_fixture): + conn, inverter_class, _ = inverters_fixture + rt_api = await solax.real_time_api(*conn, "", inverter_class.__name__) + assert rt_api.inverter.__class__ == inverter_class + assert inverter_class.__name__ in solax.discovery.get_models() + + +@pytest.mark.asyncio +async def test_discovery_unsupported_inverter(): + with pytest.raises(DiscoveryError): + await solax.real_time_api("localhost", 2, "", "doesnotexist") + + @pytest.mark.asyncio async def test_discovery_no_host(): with pytest.raises(DiscoveryError):