From 7efc79ffe37225de672693cc5e53317f91fd4bee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Thu, 23 Nov 2023 15:13:40 -0800
Subject: [PATCH] Add NPSClient.get()

---
 src/lvmnps/nps/core.py | 10 ++++++++++
 tests/test_nps.py      | 12 +++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/lvmnps/nps/core.py b/src/lvmnps/nps/core.py
index 1fd9a36..87991ed 100644
--- a/src/lvmnps/nps/core.py
+++ b/src/lvmnps/nps/core.py
@@ -105,6 +105,16 @@ async def refresh(self):
 
         pass
 
+    def get(self, outlet: int | str):
+        """Retrieves an outlet by ID or name."""
+
+        if isinstance(outlet, int):
+            return get_outlet_by_id(self.outlets, outlet)
+        elif isinstance(outlet, str):
+            return get_outlet_by_name(self.outlets, outlet)
+        else:
+            raise TypeError("Invalid outlet type. Only int and str are allowed.")
+
     async def set_state(
         self,
         outlets: OutletArgType,
diff --git a/tests/test_nps.py b/tests/test_nps.py
index 1917826..0a5b1ec 100644
--- a/tests/test_nps.py
+++ b/tests/test_nps.py
@@ -10,7 +10,7 @@
 
 import pytest
 
-from lvmnps.nps.core import NPSClient
+from lvmnps.nps.core import NPSClient, OutletModel
 
 
 async def test_npsclient(nps_test_client: NPSClient):
@@ -69,3 +69,13 @@ async def test_outlet_cycle(nps_test_client: NPSClient):
     await nps_test_client.cycle(outlet, delay=0.1)
 
     assert outlet.state is True
+
+
+@pytest.mark.parametrize("outlet", [1, "test_1"])
+async def test_client_get(nps_test_client: NPSClient, outlet: str | int):
+    assert isinstance(nps_test_client.get(outlet), OutletModel)
+
+
+async def test_client_get_invalid_type(nps_test_client: NPSClient):
+    with pytest.raises(TypeError):
+        nps_test_client.get(None)  # type: ignore