From fa36cd409616a35ee6ec5f84c798b7bb5270126b Mon Sep 17 00:00:00 2001
From: Darren Burns <darrenb900@gmail.com>
Date: Mon, 16 Oct 2023 14:39:54 +0100
Subject: [PATCH] Ensure active tab is set correctly

---
 src/textual/await_complete.py |  8 ++------
 src/textual/widgets/_tabs.py  | 33 ++++++++++++++++++++++-----------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/textual/await_complete.py b/src/textual/await_complete.py
index 078d53ecb5..8e511a9795 100644
--- a/src/textual/await_complete.py
+++ b/src/textual/await_complete.py
@@ -3,8 +3,6 @@
 from asyncio import Future, gather, wait
 from typing import Coroutine
 
-from textual._asyncio import create_task
-
 
 class AwaitComplete:
     """An 'optionally-awaitable' object.
@@ -19,16 +17,14 @@ class AwaitComplete:
 
     def __init__(self, *coroutine: Coroutine) -> None:
         self.coroutine = coroutine
-        self._future: Future | None = None
         AwaitComplete._instances.append(self)
+        self._future: Future = gather(*[coroutine for coroutine in self.coroutine])
+        self._future.add_done_callback(self._on_done)
 
     async def __call__(self):
         await self
 
     def __await__(self):
-        if not self._future:
-            self._future = gather(*self.coroutine)
-            self._future.add_done_callback(self._on_done)
         return self._future.__await__()
 
     def _on_done(self, _: Future) -> None:
diff --git a/src/textual/widgets/_tabs.py b/src/textual/widgets/_tabs.py
index 74cf7f0a52..8f37c889aa 100644
--- a/src/textual/widgets/_tabs.py
+++ b/src/textual/widgets/_tabs.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+from asyncio import create_task
 from dataclasses import dataclass
 from typing import ClassVar
 
@@ -14,7 +16,7 @@
 from ..binding import Binding, BindingType
 from ..containers import Container, Horizontal, Vertical
 from ..css.query import NoMatches
-from ..events import Mount
+from ..events import Event, Mount
 from ..geometry import Offset
 from ..message import Message
 from ..reactive import reactive
@@ -464,7 +466,7 @@ async def refresh_active() -> None:
 
         return AwaitComplete(mount_await())
 
-    def clear(self) -> AwaitRemove:
+    def clear(self) -> AwaitComplete:
         """Clear all the tabs.
 
         Returns:
@@ -474,7 +476,8 @@ def clear(self) -> AwaitRemove:
         underline.highlight_start = 0
         underline.highlight_end = 0
         self.call_after_refresh(self.post_message, self.Cleared(self))
-        return self.query("#tabs-list > Tab").remove()
+        self.active = ""
+        return AwaitComplete(self.query("#tabs-list > Tab").remove()())
 
     def remove_tab(self, tab_or_id: Tab | str | None) -> AwaitComplete:
         """Remove a tab.
@@ -498,7 +501,6 @@ def remove_tab(self, tab_or_id: Tab | str | None) -> AwaitComplete:
 
         removing_active_tab = remove_tab.has_class("-active")
         next_tab = self._next_active
-        print(f"next tab is {next_tab!r}")
         result_message: Tabs.Cleared | Tabs.TabActivated | None = None
         if removing_active_tab and next_tab is not None:
             result_message = self.TabActivated(self, next_tab)
@@ -510,13 +512,22 @@ def remove_tab(self, tab_or_id: Tab | str | None) -> AwaitComplete:
         async def do_remove() -> None:
             """Perform the remove after refresh so the underline bar gets new positions."""
             await remove_await
-            if removing_active_tab:
-                print(f"inside callback next_tab is {next_tab!r}")
-                if next_tab is not None:
-                    next_tab.add_class("-active")
-                self._highlight_active(animate=True)
-            if result_message is not None:
-                self.post_message(result_message)
+
+            removal_complete = asyncio.Event()
+
+            async def remove_after_refresh():
+                if removing_active_tab:
+                    if next_tab is not None:
+                        next_tab.add_class("-active")
+                        self.active = next_tab.id
+                    self._highlight_active(animate=True)
+                if result_message is not None:
+                    self.post_message(result_message)
+
+                removal_complete.set()
+
+            self.call_after_refresh(remove_after_refresh)
+            await removal_complete.wait()
 
         return AwaitComplete(do_remove())