async_helpers: make SingleInstanceTask.wait safer

SingleInstanceTask has distinct steps for creation of the object, and
starting the task.  If a different coroutine is waiting on the
SingleInstanceTask, it isn't safe to directly call
SingleInstanceTask.wait() as the task may or may not have been created
yet.

Existing code usage of SingleInstanceTask is in 4 categories, with
reguards to SingleInstanceTask.wait():
1) using SingleInstanceTask without using SingleInstanceTask.wait().
   This is unchanged.
2) using SingleInstanceTask.wait without a check on task is not None.
   This may be safe now, but is fragile in the face of innocent-looking
   refactors around the SingleInstanceTask.
3) using SingleInstanceTask.wait after confirming that the task is not
   None.  This is fine but a leaky abstraction.
4) directly waiting on the SingleInstanceTask.task.  Another leaky
   abstraction, but it's solving a cancellation problem.  Leaving this
   alone.

By enhancing SingleInstanceTask.wait(), cases 2 and 3 are improved.  The
code not checking the task today is made safer, and the code checking
the task today can be simplified.

subiquity/server/controllers/tests/test_filesystem.py

@@ -14,7 +14,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import copy
-from unittest import mock, TestCase, IsolatedAsyncioTestCase
+from unittest import mock, IsolatedAsyncioTestCase
 import uuid
 
 from subiquitycore.tests.parameterized import parameterized
@@ -75,7 +75,7 @@ class TestSubiquityControllerFilesystem(IsolatedAsyncioTestCase):
         self.assertTrue({'defaults', 'os'} <= actual)
 
 
-class TestGuided(TestCase):
+class TestGuided(IsolatedAsyncioTestCase):
     boot_expectations = [
         (Bootloader.UEFI, 'gpt', '/boot/efi'),
         (Bootloader.UEFI, 'msdos', '/boot/efi'),
@@ -96,7 +96,7 @@ class TestGuided(TestCase):
         self.d1 = make_disk(self.model, ptable=ptable)
 
     @parameterized.expand(boot_expectations)
-    def test_guided_direct(self, bootloader, ptable, p1mnt):
+    async def test_guided_direct(self, bootloader, ptable, p1mnt):
         self._guided_setup(bootloader, ptable)
         target = GuidedStorageTargetReformat(disk_id=self.d1.id)
         self.controller.guided(GuidedChoiceV2(target=target, use_lvm=False))
@@ -107,7 +107,7 @@ class TestGuided(TestCase):
         self.assertFalse(d1p2.preserve)
         self.assertIsNone(gaps.largest_gap(self.d1))
 
-    def test_guided_direct_BIOS_MSDOS(self):
+    async def test_guided_direct_BIOS_MSDOS(self):
         self._guided_setup(Bootloader.BIOS, 'msdos')
         target = GuidedStorageTargetReformat(disk_id=self.d1.id)
         self.controller.guided(GuidedChoiceV2(target=target, use_lvm=False))
@@ -117,7 +117,7 @@ class TestGuided(TestCase):
         self.assertIsNone(gaps.largest_gap(self.d1))
 
     @parameterized.expand(boot_expectations)
-    def test_guided_lvm(self, bootloader, ptable, p1mnt):
+    async def test_guided_lvm(self, bootloader, ptable, p1mnt):
         self._guided_setup(bootloader, ptable)
         target = GuidedStorageTargetReformat(disk_id=self.d1.id)
         self.controller.guided(GuidedChoiceV2(target=target, use_lvm=True))
@@ -133,7 +133,7 @@ class TestGuided(TestCase):
         self.assertEqual(d1p3, part)
         self.assertIsNone(gaps.largest_gap(self.d1))
 
-    def test_guided_lvm_BIOS_MSDOS(self):
+    async def test_guided_lvm_BIOS_MSDOS(self):
         self._guided_setup(Bootloader.BIOS, 'msdos')
         target = GuidedStorageTargetReformat(disk_id=self.d1.id)
         self.controller.guided(GuidedChoiceV2(target=target, use_lvm=True))
@@ -178,7 +178,7 @@ class TestGuided(TestCase):
              ('gpt', None)
          )]
     )
-    def test_guided_direct_side_by_side(self, bl, pt, flag):
+    async def test_guided_direct_side_by_side(self, bl, pt, flag):
         self._guided_side_by_side(bl, pt)
         parts_before = self.d1._partitions.copy()
         gap = gaps.largest_gap(self.d1)
@@ -198,7 +198,7 @@ class TestGuided(TestCase):
              ('gpt', None)
          )]
     )
-    def test_guided_lvm_side_by_side(self, bl, pt, flag):
+    async def test_guided_lvm_side_by_side(self, bl, pt, flag):
         self._guided_side_by_side(bl, pt)
         parts_before = self.d1._partitions.copy()
         gap = gaps.largest_gap(self.d1)
@@ -213,18 +213,18 @@ class TestGuided(TestCase):
         self.assertEqual(flag, p_data.flag)
 
 
-class TestLayout(TestCase):
+class TestLayout(IsolatedAsyncioTestCase):
     def setUp(self):
         self.app = make_app()
         self.app.opts.bootloader = None
         self.fsc = FilesystemController(app=self.app)
 
     @parameterized.expand([('reformat_disk',), ('use_gap',)])
-    def test_good_modes(self, mode):
+    async def test_good_modes(self, mode):
         self.fsc.validate_layout_mode(mode)
 
     @parameterized.expand([('resize_biggest',), ('use_free',)])
-    def test_bad_modes(self, mode):
+    async def test_bad_modes(self, mode):
         with self.assertRaises(ValueError):
             self.fsc.validate_layout_mode(mode)
 

subiquitycore/async_helpers.py

@@ -73,6 +73,7 @@ class SingleInstanceTask:
     def __init__(self, func, propagate_errors=True, cancel_restart=True):
         self.func = func
         self.propagate_errors = propagate_errors
+        self.task_created = asyncio.Event()
         self.task = None
         # if True, allow subsequent start calls to cancel a running task
         # raises TaskAlreadyRunningError if we skip starting the task.
@@ -102,9 +103,11 @@ class SingleInstanceTask:
             self.task = asyncio.Task(coro)
         else:
             self.task = coro
+        self.task_created.set()
         return schedule_task(self._start(old))
 
     async def wait(self):
+        await self.task_created.wait()
         while True:
             try:
                 return await self.task

subiquitycore/tests/test_async_helpers.py

@@ -45,3 +45,21 @@ class TestSingleInstanceTask(unittest.IsolatedAsyncioTestCase):
         sit.task.cancel()
         self.assertEqual(expected_call_count, mock_fn.call_count)
         self.assertEqual(cancel_restart, restarted)
+
+
+# previously, wait() may or may not have been safe to call, depending
+# on if the task had actually been created yet.
+class TestSITWait(unittest.IsolatedAsyncioTestCase):
+    async def test_wait_started(self):
+        async def fn():
+            pass
+        sit = SingleInstanceTask(fn)
+        await sit.start()
+        await asyncio.wait_for(sit.wait(), timeout=1.0)
+
+    async def test_wait_not_started(self):
+        async def fn():
+            self.fail('not supposed to be called')
+        sit = SingleInstanceTask(fn)
+        with self.assertRaises(asyncio.TimeoutError):
+            await asyncio.wait_for(sit.wait(), timeout=0.1)