From 4995f576de09dc5a521e8ce9304c94ab9f4c4062 Mon Sep 17 00:00:00 2001
From: Jesse Schoepfer <jelloeater@gmail.com>
Date: Mon, 22 May 2023 02:43:49 -0400
Subject: [PATCH] wip: Fixing CLI test

---
 hubitatcontrol/__main__.py | 45 ++++++++++----------------------------
 tests/test_cli.py          |  2 +-
 tests/test_main.py         |  3 +--
 3 files changed, 13 insertions(+), 37 deletions(-)

diff --git a/hubitatcontrol/__main__.py b/hubitatcontrol/__main__.py
index 2ca0f67..d781953 100644
--- a/hubitatcontrol/__main__.py
+++ b/hubitatcontrol/__main__.py
@@ -7,6 +7,7 @@
 import typer
 from dotenv import load_dotenv
 
+import hubitatcontrol
 from hubitatcontrol import Hub
 
 app_name = "hubitatcontrol"
@@ -41,25 +42,6 @@ def hub_from_keyring():
     return get_hub(host=host_env, token=token_env, app_id=app_id_env, cloud_token=cloud_token_env)
 
 
-def lookup_device(hub_in, device_lookup):
-    """
-    Takes device NAME, not ID for lookup
-    """
-
-    return get_device_type(
-        device_in=hub_in.get_device(device_lookup), hub_in=hub_in
-    )  # Fall through return # pragma: no cover
-
-
-def lookup_device_id(hub_in, device_id):
-    """
-    Takes device ID for lookup
-    """
-    return get_device_type(
-        device_in=hub_in.get_device_id(device_id), hub_in=hub_in
-    )  # Fall through return # pragma: no cover
-
-
 @app.command()
 def ls():
     """
@@ -70,18 +52,17 @@ def ls():
 
 
 @app.command()
-def on(device_id: int):
+def on(device_id: int) -> hubitatcontrol.lights.Switch:
     """
     Turn on a device via it's Device ID
     """
 
     hub_in = hub_from_keyring()
-    device = hub_in.get_device_id(device_id)
-    import hubitatcontrol.hub
+    dev = hubitatcontrol.GetDevices(hub_in).Switch()
 
-    dev = lookup_device(hub_in, device['name'])
-
-    dev.turn_on()
+    for i in dev:
+        if i.id == device_id:
+            i.turn_on()
 
 
 @app.command()
@@ -91,12 +72,11 @@ def off(device_id: int):
     """
 
     hub_in = hub_from_keyring()
-    device = hub_in.get_device_id(device_id)
-    import hubitatcontrol.hub
-
-    dev = lookup_device(hub_in, device['name'])
+    dev = hubitatcontrol.GetDevices(hub_in).Switch()
 
-    dev.turn_off()
+    for i in dev:
+        if i.id == device_id:
+            i.turn_off()
 
 
 @app.command()
@@ -104,10 +84,7 @@ def level(device_id: int, level: int):
     """Turn on a device via it's Device ID"""
 
     hub_in = hub_from_keyring()
-    device = hub_in.get_device_id(device_id)
-    import hubitatcontrol.hub
-
-    dev = lookup_device(hub_in, device['name'])
+    dev = hub_in.get_device_id(device_id)
 
     dev.set_level(level)
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index e0cf75e..5323c6c 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -35,4 +35,4 @@ def test_cli_keyring(self):
     @pytest.mark.skipif(load_dotenv('.env') is False, reason='No env file found')
     def test_cli_on(self):
         test_bulb = get_device_of_type("Virtual RGBW Light")
-        cli.on(test_bulb.id)
+        cli.on(test_bulb['id'])
diff --git a/tests/test_main.py b/tests/test_main.py
index 527603a..f1a25a9 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -2,7 +2,6 @@
 import random
 import time
 
-from __main__ import lookup_device
 from dotenv import load_dotenv
 
 import hubitatcontrol
@@ -19,7 +18,7 @@ def get_device_of_type(device_type: str):
     h = Hub(host=host_env, token=token_env, app_id=app_id_env, cloud_token=cloud_token)
     for i in h.devices:
         if i["type"] == device_type:
-            return lookup_device(h, i["label"])
+            return i
 
 
 def test_creds():