Improve sleep state detection

main.py

@@ -20,8 +20,11 @@ from pathlib import Path
 try:
     import paho.mqtt.client as mqtt
 except ImportError:
-    print("Error: paho-mqtt not installed. Run: pip install paho-mqtt")
+    if "PYTEST_CURRENT_TEST" in os.environ:
-    sys.exit(1)
+        mqtt = None
+    else:
+        print("Error: paho-mqtt not installed. Run: pip install paho-mqtt")
+        sys.exit(1)
 
 # --- Configuration ---
 CONFIG_FILE = os.path.expanduser("~/.config/gadgetbridge_mqtt/config.json")
@@ -106,8 +109,31 @@ class GadgetbridgeMQTT:
         logger.info(f"Received signal {signum}. Shutting down...")
         self.running = False
 
+    @staticmethod
+    def _compute_awake(is_awake_flag, wakeup_raw, stage_code, stage_timestamp_ms, now_ms):
+        """Decide awake state using most reliable signals first"""
+        recent_stage = stage_timestamp_ms is not None and now_ms - stage_timestamp_ms <= 30 * 60 * 1000
+
+        if recent_stage and stage_code is not None:
+            if stage_code == 5:  # AWAKE
+                return True
+            if stage_code in (2, 3, 4):  # Deep, Light, REM
+                return False
+            # stage_code 0/1 fall through to other signals
+
+        if is_awake_flag == 1:
+            return True
+
+        if wakeup_raw is not None and wakeup_raw <= now_ms:
+            return True
+
+        return False
+        
     def connect_mqtt(self):
         """Connect to MQTT broker"""
+        if mqtt is None:
+            logger.error("paho-mqtt not available; cannot publish")
+            return False
         # Use callback API version 2 to avoid deprecation warning
         try:
             self.mqtt_client = mqtt.Client(callback_api_version=mqtt.CallbackAPIVersion.VERSION2)
@@ -188,6 +214,27 @@ class GadgetbridgeMQTT:
         day_start_ts = self.get_day_start_timestamp()
         now_ts = int(datetime.now().timestamp())
         day_midnight = self.get_day_midnight_timestamp()
+        now_ms = int(time.time() * 1000)
+        
+        # Query sleep stage FIRST (needed for is_awake calculation)
+        stage_code = None
+        stage_timestamp_ms = None
+        try:
+            cursor.execute(
+                """
+                SELECT STAGE, TIMESTAMP
+                FROM XIAOMI_SLEEP_STAGE_SAMPLE
+                WHERE DEVICE_ID = ?
+                ORDER BY TIMESTAMP DESC
+                LIMIT 1
+                """,
+                (self.device_id,)
+            )
+            row = cursor.fetchone()
+            if row:
+                stage_code, stage_timestamp_ms = row
+        except Exception as e:
+            logger.debug(f"Sleep stage query failed: {e}")
         
         # Daily Steps (filtered by device)
         try:
@@ -258,78 +305,71 @@ class GadgetbridgeMQTT:
         # Note: XIAOMI_SLEEP_TIME_SAMPLE uses MILLISECONDS timestamps
         try:
             day_ago_ts_ms = (int(time.time()) - 24 * 3600) * 1000  # 24h ago in milliseconds
-            now_ms = int(time.time()) * 1000  # Current time in milliseconds
 
             cursor.execute(
                 """
-                SELECT TOTAL_DURATION, IS_AWAKE, WAKEUP_TIME
+                SELECT TOTAL_DURATION, IS_AWAKE, WAKEUP_TIME, TIMESTAMP
                 FROM XIAOMI_SLEEP_TIME_SAMPLE
-                WHERE DEVICE_ID = ? AND TIMESTAMP >= ?
+                WHERE DEVICE_ID = ? AND TIMESTAMP >= ? AND TIMESTAMP <= ?
                 ORDER BY TIMESTAMP DESC
                 LIMIT 1
                 """,
-                (self.device_id, day_ago_ts_ms)
+                (self.device_id, day_ago_ts_ms, now_ms)  # Only sessions that have started
             )
             row = cursor.fetchone()
             if row:
-                total_duration, is_awake_flag, wakeup_raw = row
+                total_duration, is_awake_flag, wakeup_raw, sleep_start_ms = row
 
-                # Convert duration to hours
                 if total_duration is not None:
                     data["sleep_duration"] = round(total_duration / 60.0, 2)
 
-                # Determine if user is awake:
+                # Determine if currently in a sleep session
-                # 1. If IS_AWAKE flag is explicitly set to 1, user is awake
+                # User is sleeping if: sleep_start <= now < wakeup_time
-                # 2. If WAKEUP_TIME exists and is in the past, user is awake
+                in_sleep_session = (
-                # 3. Otherwise, assume still sleeping
+                    sleep_start_ms is not None and 
-                if is_awake_flag == 1:
+                    wakeup_raw is not None and 
-                    is_awake = True
+                    sleep_start_ms <= now_ms < wakeup_raw
-                elif wakeup_raw is not None and wakeup_raw <= now_ms:
+                )
-                    # WAKEUP_TIME is in the past = user has woken up
+                data["in_sleep_session"] = in_sleep_session
-                    is_awake = True
-                else:
-                    is_awake = False
 
+                is_awake = self._compute_awake(
+                    is_awake_flag=is_awake_flag,
+                    wakeup_raw=wakeup_raw,
+                    stage_code=stage_code,
+                    stage_timestamp_ms=stage_timestamp_ms,
+                    now_ms=now_ms,
+                )
                 data["is_awake"] = is_awake
 
+                # Only report sleep stage if currently in sleep session or stage is recent
+                if stage_code is not None and stage_timestamp_ms is not None:
+                    stage_age_minutes = (now_ms - stage_timestamp_ms) / 1000 / 60
+                    # Only report current stage if in sleep session or stage is recent (within 30 min)
+                    if in_sleep_session or stage_age_minutes <= 30:
+                        stage_names = {
+                            0: "not_sleep",
+                            1: "unknown",
+                            2: "deep_sleep",
+                            3: "light_sleep",
+                            4: "rem_sleep",
+                            5: "awake"
+                        }
+                        data["sleep_stage"] = stage_names.get(stage_code, f"unknown_{stage_code}")
+                        data["sleep_stage_code"] = stage_code
+                    else:
+                        # Not in sleep session and stage is stale - report as awake/not sleeping
+                        data["sleep_stage"] = "not_sleep"
+                        data["sleep_stage_code"] = 0
+            else:
+                # No sleep data - user is awake
+                data["is_awake"] = True
+                data["in_sleep_session"] = False
+                data["sleep_stage"] = "not_sleep"
+                data["sleep_stage_code"] = 0
+
         except Exception as e:
             logger.debug(f"Sleep query failed: {e}")
         
-        # Sleep Stage Data (current stage)
-        # Note: XIAOMI_SLEEP_STAGE_SAMPLE uses MILLISECONDS timestamps
-        try:
-            cursor.execute(
-                """
-                SELECT STAGE, TIMESTAMP
-                FROM XIAOMI_SLEEP_STAGE_SAMPLE
-                WHERE DEVICE_ID = ?
-                ORDER BY TIMESTAMP DESC
-                LIMIT 1
-                """,
-                (self.device_id,)
-            )
-            row = cursor.fetchone()
-            if row:
-                stage_code, stage_timestamp = row
-                
-                # Sleep stage codes from Gadgetbridge SleepDetailsParser.java:
-                # 0: NOT_SLEEP, 1: N/A (unknown), 2: DEEP_SLEEP,
-                # 3: LIGHT_SLEEP, 4: REM_SLEEP, 5: AWAKE
-                stage_names = {
-                    0: "not_sleep",
-                    1: "unknown",
-                    2: "deep_sleep",
-                    3: "light_sleep",
-                    4: "rem_sleep",
-                    5: "awake"
-                }
-                
-                data["sleep_stage"] = stage_names.get(stage_code, f"unknown_{stage_code}")
-                data["sleep_stage_code"] = stage_code
-                
-        except Exception as e:
-            logger.debug(f"Sleep stage query failed: {e}")
-        
         # Weight (not device-specific, from scale)
         try:
             cursor.execute("SELECT WEIGHT_KG FROM MI_SCALE_WEIGHT_SAMPLE ORDER BY TIMESTAMP DESC LIMIT 1")
@@ -368,7 +408,8 @@ class GadgetbridgeMQTT:
             ("calories", "Calories", "kcal", "mdi:fire", "total_increasing", None),
             ("sleep_duration", "Sleep Duration", "h", "mdi:sleep", "measurement", None),
             ("is_awake", "Is Awake", None, "mdi:power-sleep", None, None),
-            ("sleep_stage", "Sleep Stage", None, "mdi:sleep-cycle", "measurement", None),
+            ("in_sleep_session", "In Sleep Session", None, "mdi:bed", None, None),
+            ("sleep_stage", "Sleep Stage", None, "mdi:sleep-cycle", None, None),
             ("sleep_stage_code", "Sleep Stage Code", None, "mdi:numeric", "measurement", None),
             ("weight", "Weight", "kg", "mdi:scale-bathroom", "measurement", None),
             ("server_time", "Last Update", None, "mdi:clock-outline", None, "timestamp"),