1 files changed, 59 insertions, 0 deletions
diff --git a/t/test_functionality.py b/t/test_functionality.py
index b60cb46..63863cf 100644
--- a/t/test_functionality.py
+++ b/t/test_functionality.py
@@ -168,3 +168,62 @@ class BasicUsage(utils.AergiaUnitTestCase):
         self.assertFuncContains('c', [self.expected_samples(delay)], samples)
         self.assertFuncContains('b', [], samples)
         self.assertFuncContains('a', [], samples)
+
+    def test_task_groups(self):
+        d1 = 0.2
+        d2 = 0.3
+        d3 = 0.1
+        async def d(): await asyncio.sleep(d1)
+        async def c(): await asyncio.sleep(d2)
+        async def b(): await asyncio.sleep(d3)
+
+        async def a():
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(d())
+                tg.create_task(c())
+                tg.create_task(b())
+
+        self.Aergia.start()
+        asyncio.run(a())
+        self.Aergia.stop()
+
+        samples = self.Aergia.get_samples()
+
+        self.assertFuncContains('d', [self.expected_samples(d1)], samples)
+        self.assertFuncContains('c', [self.expected_samples(d2)], samples)
+        self.assertFuncContains('b', [self.expected_samples(d3)], samples)
+        # TODO where does this come from?
+        self.assertFuncContains('a', [self.expected_samples(d2)], samples)
+
+    def test_task_groups_cancel(self):
+        d1 = 0.1
+        d2 = 0.2
+        d3 = 0.3
+        async def d(): await asyncio.sleep(d1)
+
+        async def c():
+            await asyncio.sleep(d2)
+            1 / 0  # crash
+
+        async def b(): await asyncio.sleep(d3)
+
+        async def a():
+            try:
+                async with asyncio.TaskGroup() as tg:
+                    tg.create_task(d())
+                    tg.create_task(c())
+                    tg.create_task(b())
+            except:
+                pass
+
+        self.Aergia.start()
+        asyncio.run(a())
+        self.Aergia.stop()
+
+        samples = self.Aergia.get_samples()
+
+        self.assertFuncContains('d', [self.expected_samples(d1)], samples)
+        self.assertFuncContains('c', [self.expected_samples(d2)], samples)
+        self.assertFuncContains('b', [self.expected_samples(d2)], samples)
+        # TODO where does this come from?
+        self.assertFuncContains('a', [self.expected_samples(d2)], samples)