From df5df4b497cd9263563677f117568227577c960b Mon Sep 17 00:00:00 2001
From: Burak Ozter <burak.ozter@spirirobotics.com>
Date: Thu, 7 Nov 2024 13:54:04 -0400
Subject: [PATCH] FIX: button state when stop/delete and video stream was
 enabled

---
 guiTools/spiri_sdk_guitools/sim_drone.py    |  9 ++++++++-
 guiTools/spiri_sdk_guitools/video_button.py | 15 +++++++++++----
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/guiTools/spiri_sdk_guitools/sim_drone.py b/guiTools/spiri_sdk_guitools/sim_drone.py
index 711d2b0..3cc29b9 100644
--- a/guiTools/spiri_sdk_guitools/sim_drone.py
+++ b/guiTools/spiri_sdk_guitools/sim_drone.py
@@ -85,6 +85,7 @@ class Robot:
             ]
         self.compose_files = compose_files
         self.processes = []
+        self.video_button = None
         robots.add(self)
 
     async def ui_containers(self, element):
@@ -142,7 +143,9 @@ class Robot:
                 ui.label(f"""Sysid: {self.sysid}""")
                 ui.button("Start", on_click=self.start).classes("m-2")
                 ui.button("Stop", on_click=self.async_stop).classes("m-2")
-                EnableStreamingButton(sysid=self.sysid).classes("m-2")
+                self.video_button = EnableStreamingButton(sysid=self.sysid).classes(
+                    "m-2"
+                )
 
                 async def delete_robot():
                     await self.async_stop()
@@ -164,6 +167,10 @@ class Robot:
         return await run.io_bound(self.stop)
 
     def stop(self):
+        # If video was enabled, and we stop. Button state is wrong.
+        # Stop video if it is enabled.
+        if isinstance(self.video_button, EnableStreamingButton):
+            self.video_button.stop_video()
         # Delete gazebo model
         self.delete_gz_model(sysid=self.sysid)
         # Signal all processes to stop
diff --git a/guiTools/spiri_sdk_guitools/video_button.py b/guiTools/spiri_sdk_guitools/video_button.py
index 07205a0..7c29164 100644
--- a/guiTools/spiri_sdk_guitools/video_button.py
+++ b/guiTools/spiri_sdk_guitools/video_button.py
@@ -27,18 +27,25 @@ class EnableStreamingButton(ui.element):
         result = await run.cpu_bound(self.enable_streaming, self.sysid, not self._state)
         if result:
             ui.notify("Success", type="positive]")
-            self._state = not self._state
-            self.update()
+            self.set_state(state=not self._state)
         else:
-            self._state = False
-            self.update()
+            self.set_state(state=False)
             ui.notify("No Video Streaming Available..", type="negative")
         spinner.delete()
 
+    def set_state(self, state: bool):
+        self._state = state
+        self.update()
+
     def update(self) -> None:
         self.button._text = f'{"Disable" if self._state else "Enable"}' + " Video"
         self.button.props(f'color={"red" if self._state else "green"}')
 
+    def stop_video(self):
+        if self._state != False:
+            self.enable_streaming(sysid=self.sysid, is_streaming=False)
+            self.set_state(state=False)
+
     @staticmethod
     def enable_streaming(sysid: int, is_streaming: bool) -> bool: