diff --git a/smarttub/api.py b/smarttub/api.py
index aa639b9..09b10ea 100644
--- a/smarttub/api.py
+++ b/smarttub/api.py
@@ -196,8 +196,10 @@ async def _wait_for_state_change(
             RuntimeError if the state change is not reflected within the timeout period
         """
         start_time = datetime.datetime.now().timestamp()
+        # Use the provided method if available, otherwise use default get_status
+        status_method = get_status_method if get_status_method else self.get_status
         while True:
-            state = await self.get_status()
+            state = await status_method()
             if check_func(state):
                 return state
 
@@ -206,9 +208,6 @@ async def _wait_for_state_change(
 
             await asyncio.sleep(0.5)
 
-            if get_status_method:
-                state = await get_status_method()
-
     async def get_status(self) -> "SpaState":
         """Query the status of the spa."""
         return SpaState(self, **await self.request("GET", "status"))
@@ -407,10 +406,12 @@ class SpaStateFull(SpaState):
     def __init__(self, spa: Spa, state: dict):
         super().__init__(spa, **state)
         self.lights = [
-            SpaLight(spa, **light_props) for light_props in self.properties["lights"]
+            SpaLight(spa, **light_props)
+            for light_props in (self.properties.get("lights") or [])
         ]
         self.pumps = [
-            SpaPump(spa, **pump_props) for pump_props in self.properties["pumps"]
+            SpaPump(spa, **pump_props)
+            for pump_props in (self.properties.get("pumps") or [])
         ]
         self.sensors = [
             SpaSensor(spa, **sensor_props)