From 394cfc723c95729081584e8d47915a8543038401 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Mon, 16 Dec 2019 12:41:43 +1300
Subject: [PATCH] handle check task being cancelled while view is shown

---
 subiquity/controllers/refresh.py |  4 +---
 subiquity/ui/views/refresh.py    | 16 ++++++++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/subiquity/controllers/refresh.py b/subiquity/controllers/refresh.py
index e76eaa7b..2c32e537 100644
--- a/subiquity/controllers/refresh.py
+++ b/subiquity/controllers/refresh.py
@@ -65,9 +65,7 @@ class RefreshController(BaseController):
     @property
     def check_state(self):
         task = self.check_task.task
-        if not task.done():
-            return CheckState.UNKNOWN
-        if task.cancelled():
+        if not task.done() or task.cancelled():
             return CheckState.UNKNOWN
         if task.exception():
             return CheckState.UNAVAILABLE
diff --git a/subiquity/ui/views/refresh.py b/subiquity/ui/views/refresh.py
index b4f72ce7..973b7a57 100644
--- a/subiquity/ui/views/refresh.py
+++ b/subiquity/ui/views/refresh.py
@@ -163,10 +163,18 @@ class RefreshView(BaseView):
         schedule_task(self._wait_check_result())
 
     async def _wait_check_result(self):
-        try:
-            check_state = await self.controller.check_task.task
-        except Exception as e:
-            self.check_state_failed(e)
+        while True:
+            try:
+                check_state = await self.controller.check_task.task
+            except asyncio.CancelledError:
+                # If the task was cancelled, another will have been
+                # started.
+                continue
+            except Exception as e:
+                self.check_state_failed(e)
+                return
+            else:
+                break
         if check_state == CheckState.AVAILABLE:
             self.check_state_available()
         else: