From 0e9729f601f6e69a7b19f41200d255cc50ddf0a6 Mon Sep 17 00:00:00 2001
From: Fabian Peter Hammerle <fabian@hammerle.me>
Date: Thu, 13 May 2021 08:57:43 +0200
Subject: [PATCH] curtains: report position after sending stop command

rebased https://github.com/fphammerle/switchbot-mqtt/pull/31/commits/885d557c51c70e6692810e8d3f395d495a560b4d
---
 CHANGELOG.md                          |  3 +++
 setup.py                              |  2 +-
 switchbot_mqtt/__init__.py            |  5 +++++
 tests/test_switchbot_curtain_motor.py | 23 ++++++++++++++++++++++-
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ff6d61e..a0272ce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- Report position of curtain motor on topic  `homeassistant/cover/switchbot-curtain/MAC_ADDRESS/position`
+  after sending stop command.
 
 ## [1.0.0] - 2021-07-25
 ### Added
diff --git a/setup.py b/setup.py
index 081df23..37482e0 100644
--- a/setup.py
+++ b/setup.py
@@ -72,7 +72,7 @@
     ],
     entry_points={"console_scripts": ["switchbot-mqtt = switchbot_mqtt:_main"]},
     install_requires=[
-        # >=0.10.0 for SwitchbotCurtain.get_position
+        # >=0.10.0 for SwitchbotCurtain.{update,get_position}
         "PySwitchbot>=0.10.0,<0.11",
         "paho-mqtt<2",
     ],
diff --git a/switchbot_mqtt/__init__.py b/switchbot_mqtt/__init__.py
index 0627bbd..eaf261d 100644
--- a/switchbot_mqtt/__init__.py
+++ b/switchbot_mqtt/__init__.py
@@ -265,6 +265,10 @@ def _report_position(self, mqtt_client: paho.mqtt.client.Client) -> None:
             mqtt_client=mqtt_client,
         )
 
+    def _update_position(self, mqtt_client: paho.mqtt.client.Client) -> None:
+        self._device.update()
+        self._report_position(mqtt_client=mqtt_client)
+
     def execute_command(
         self, mqtt_message_payload: bytes, mqtt_client: paho.mqtt.client.Client
     ) -> None:
@@ -293,6 +297,7 @@ def execute_command(
                 # https://www.home-assistant.io/integrations/cover.mqtt/#configuration-variables
                 # https://community.home-assistant.io/t/mqtt-how-to-remove-retained-messages/79029/2
                 self.report_state(mqtt_client=mqtt_client, state=b"")
+                self._update_position(mqtt_client=mqtt_client)
         else:
             _LOGGER.warning(
                 "unexpected payload %r (expected 'OPEN', 'CLOSE', or 'STOP')",
diff --git a/tests/test_switchbot_curtain_motor.py b/tests/test_switchbot_curtain_motor.py
index 38c906e..214d184 100644
--- a/tests/test_switchbot_curtain_motor.py
+++ b/tests/test_switchbot_curtain_motor.py
@@ -94,6 +94,21 @@ def test__report_position_invalid(caplog, position):
     publish_mock.assert_not_called()
 
 
+def test__update_position():
+    with unittest.mock.patch("switchbot.SwitchbotCurtain.__init__", return_value=None):
+        actor = switchbot_mqtt._CurtainMotor(
+            mac_address="dummy", retry_count=21, password=None
+        )
+    with unittest.mock.patch(
+        "switchbot.SwitchbotCurtain.update"
+    ) as update_mock, unittest.mock.patch.object(
+        actor, "_report_position"
+    ) as report_position_mock:
+        actor._update_position(mqtt_client="client")
+    update_mock.assert_called_once_with()
+    report_position_mock.assert_called_once_with(mqtt_client="client")
+
+
 @pytest.mark.parametrize("mac_address", ["aa:bb:cc:dd:ee:ff", "aa:bb:cc:11:22:33"])
 @pytest.mark.parametrize("password", ["pa$$word", None])
 @pytest.mark.parametrize("retry_count", (2, 3))
@@ -131,7 +146,9 @@ def test_execute_command(
             actor, "report_state"
         ) as report_mock, unittest.mock.patch(
             action_name, return_value=command_successful
-        ) as action_mock:
+        ) as action_mock, unittest.mock.patch.object(
+            actor, "_update_position"
+        ) as update_position_mock:
             actor.execute_command(
                 mqtt_client="dummy", mqtt_message_payload=message_payload
             )
@@ -170,6 +187,10 @@ def test_execute_command(
             )
         ]
         report_mock.assert_not_called()
+    if action_name == "switchbot.SwitchbotCurtain.stop" and command_successful:
+        update_position_mock.assert_called_once_with(mqtt_client="dummy")
+    else:
+        update_position_mock.assert_not_called()
 
 
 @pytest.mark.parametrize("mac_address", ["aa:bb:cc:dd:ee:ff"])