diff --git a/__init__.py b/__init__.py
index 8595f13..314c9a1 100644
--- a/__init__.py
+++ b/__init__.py
@@ -30,7 +30,7 @@ from .const import (
     ERRORS,
     MIN_TIME_BETWEEN_UPDATES,
     MODE,
-    ROBOT_ACTION_DOCKING,
+    ROBOT_CLEANING_ACTIONS,
     ROBOT_STATE_BUSY,
     ROBOT_STATE_ERROR,
     ROBOT_STATE_IDLE,
@@ -186,16 +186,17 @@ class VorwerkState:
         self._update_robot_info()
         self._update_state()
 
-    def _update_state(self):
+    def _update_robot_info(self):
         try:
             if not self.robot_info:
                 self.robot_info = self.robot.get_general_info().json().get("data")
         except NeatoRobotException:
             _LOGGER.warning("Couldn't fetch robot information of %s", self.robot.name)
 
-    def _update_robot_info(self):
+    def _update_state(self):
         try:
             self.robot_state = self.robot.state
+            _LOGGER.debug(self.robot_state)
         except NeatoRobotException as ex:
             if self.available:  # print only once when available
                 _LOGGER.error(
@@ -236,10 +237,11 @@ class VorwerkState:
         elif robot_state == ROBOT_STATE_IDLE:
             state = STATE_IDLE
         elif robot_state == ROBOT_STATE_BUSY:
-            if robot_state["action"] != ROBOT_ACTION_DOCKING:
-                state = STATE_RETURNING
-            else:
+            action = self.robot_state.get("action")
+            if action in ROBOT_CLEANING_ACTIONS:
                 state = STATE_CLEANING
+            else:
+                state = STATE_RETURNING
         elif robot_state == ROBOT_STATE_PAUSE:
             state = STATE_PAUSED
         elif robot_state == ROBOT_STATE_ERROR:
@@ -282,21 +284,19 @@ class VorwerkState:
 
     def _error_status(self):
         """Return error status."""
-        robot_state = self.robot_state.get("state")
-        return ERRORS.get(robot_state["error"], robot_state["error"])
+        return ERRORS.get(self.robot_state["error"], self.robot_state["error"])
 
     def _cleaning_status(self):
         """Return cleaning status."""
-        robot_state = self.robot_state.get("state")
         status_items = [
-            MODE.get(robot_state["cleaning"]["mode"]),
-            ACTION.get(robot_state["action"]),
+            MODE.get(self.robot_state["cleaning"]["mode"]),
+            ACTION.get(self.robot_state["action"]),
         ]
         if (
-            "boundary" in robot_state["cleaning"]
-            and "name" in robot_state["cleaning"]["boundary"]
+            "boundary" in self.robot_state["cleaning"]
+            and "name" in self.robot_state["cleaning"]["boundary"]
         ):
-            status_items.append(robot_state["cleaning"]["boundary"]["name"])
+            status_items.append(self.robot_state["cleaning"]["boundary"]["name"])
         return " ".join(s for s in status_items if s)
 
     @property
diff --git a/const.py b/const.py
index d5c45bf..83eba21 100644
--- a/const.py
+++ b/const.py
@@ -20,27 +20,64 @@ VORWERK_CLIENT_ID = "KY4YbVAvtgB7lp8vIbWQ7zLk3hssZlhR"
 
 MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=1)
 
-MODE = {1: "Eco", 2: "Turbo"}
+ATTR_NAVIGATION = "navigation"
+ATTR_CATEGORY = "category"
+ATTR_ZONE = "zone"
+
+ROBOT_STATE_INVALID = 0
+ROBOT_STATE_IDLE = 1
+ROBOT_STATE_BUSY = 2
+ROBOT_STATE_PAUSE = 3
+ROBOT_STATE_ERROR = 4
+
+ROBOT_ACTION_INVALID = 0
+ROBOT_ACTION_HOUSE_CLEANING = 1
+ROBOT_ACTION_SPOT_CLEANING = 2
+ROBOT_ACTION_MANUAL_CLEANING = 3
+ROBOT_ACTION_DOCKING = 4
+ROBOT_ACTION_USER_MENU = 5
+ROBOT_ACTION_SUSPENDED_CLEANING = 6
+ROBOT_ACTION_UPDATING = 7
+ROBOT_ACTION_COPY_LOGS = 8
+ROBOT_ACTION_RECOVERING_LOCATION = 9
+ROBOT_ACTION_IEC_TEST = 10
+ROBOT_ACTION_MAP_CLEANING = 11
+ROBOT_ACTION_EXPLORING_MAP = 12
+ROBOT_ACTION_ACQUIRING_MAP_IDS = 13
+ROBOT_ACTION_UPLOADING_MAP = 14
+ROBOT_ACTION_SUSPENDED_EXPLORATION = 15
+
+ROBOT_CLEANING_ACTIONS = [
+    ROBOT_ACTION_HOUSE_CLEANING,
+    ROBOT_ACTION_SPOT_CLEANING,
+    ROBOT_ACTION_MANUAL_CLEANING,
+    ROBOT_ACTION_SUSPENDED_CLEANING,
+    ROBOT_ACTION_MAP_CLEANING,
+    ROBOT_ACTION_EXPLORING_MAP,
+    ROBOT_ACTION_SUSPENDED_EXPLORATION
+]
 
 ACTION = {
-    0: "Invalid",
-    1: "House Cleaning",
-    2: "Spot Cleaning",
-    3: "Manual Cleaning",
-    4: "Docking",
-    5: "User Menu Active",
-    6: "Suspended Cleaning",
-    7: "Updating",
-    8: "Copying logs",
-    9: "Recovering Location",
-    10: "IEC test",
-    11: "Map cleaning",
-    12: "Exploring map (creating a persistent map)",
-    13: "Acquiring Persistent Map IDs",
-    14: "Creating & Uploading Map",
-    15: "Suspended Exploration",
+    ROBOT_ACTION_INVALID: "Invalid",
+    ROBOT_ACTION_HOUSE_CLEANING: "House Cleaning",
+    ROBOT_ACTION_SPOT_CLEANING: "Spot Cleaning",
+    ROBOT_ACTION_MANUAL_CLEANING: "Manual Cleaning",
+    ROBOT_ACTION_DOCKING: "Docking",
+    ROBOT_ACTION_USER_MENU: "User Menu Active",
+    ROBOT_ACTION_SUSPENDED_CLEANING: "Suspended Cleaning",
+    ROBOT_ACTION_UPDATING: "Updating",
+    ROBOT_ACTION_COPY_LOGS: "Copying logs",
+    ROBOT_ACTION_RECOVERING_LOCATION: "Recovering Location",
+    ROBOT_ACTION_IEC_TEST: "IEC test",
+    ROBOT_ACTION_MAP_CLEANING: "Map cleaning",
+    ROBOT_ACTION_EXPLORING_MAP: "Exploring map (creating a persistent map)",
+    ROBOT_ACTION_ACQUIRING_MAP_IDS: "Acquiring Persistent Map IDs",
+    ROBOT_ACTION_UPLOADING_MAP: "Creating & Uploading Map",
+    ROBOT_ACTION_SUSPENDED_EXPLORATION: "Suspended Exploration",
 }
 
+MODE = {1: "Eco", 2: "Turbo"}
+
 ERRORS = {
     "ui_error_battery_battundervoltlithiumsafety": "Replace battery",
     "ui_error_battery_critical": "Replace battery",
@@ -157,19 +194,3 @@ ALERTS = {
     "clean_incomplete_to_start": "Cleaning incomplete",
     "log_upload_failed": "Logs failed to upload",
 }
-
-ATTR_NAVIGATION = "navigation"
-ATTR_CATEGORY = "category"
-ATTR_ZONE = "zone"
-
-
-ROBOT_STATE_INVALID = 0
-ROBOT_STATE_IDLE = 1
-ROBOT_STATE_BUSY = 2
-ROBOT_STATE_PAUSE = 3
-ROBOT_STATE_ERROR = 4
-
-ROBOT_ACTION_HOUSE_CLEANING = 1
-ROBOT_ACTION_SPOT_CLEANING = 2
-ROBOT_ACTION_MANUAL_CLEANING = 3
-ROBOT_ACTION_DOCKING = 4