3 files changed, 59 insertions, 55 deletions
diff --git a/t/eager_and_scheduled.py b/t/eager_and_scheduled.py
deleted file mode 100644
index 030b539..0000000
--- a/t/eager_and_scheduled.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import asyncio
-
-
-async def foo():
-    await asyncio.sleep(1.0)
-    await baz()
-    await asyncio.sleep(0.5)
-
-
-async def bar():
-    proc = await asyncio.create_subprocess_shell('sleep 1.0')
-    await proc.communicate()
-
-
-async def baz():
-    await asyncio.sleep(1.0)
-
-
-asyncio.run(foo())
-asyncio.run(bar())
diff --git a/t/task_groups_and_cancel.py b/t/task_groups_and_cancel.py
deleted file mode 100644
index dcc6bbe..0000000
--- a/t/task_groups_and_cancel.py
+++ /dev/null
@@ -1,35 +0,0 @@
-import asyncio
-
-
-async def sleep():
-    await asyncio.sleep(3)
-    print('I should never finish!')
-    return 0
-
-
-async def work():
-    i = 0
-    while i < 50:
-        i += 1
-    await asyncio.sleep(0.2)
-    return 0
-
-
-async def explode():
-    await asyncio.sleep(1.5)
-    a = 1 / 0
-    return a
-
-
-async def main():
-    # exploding will bring all other tasks down with it!
-    try:
-        async with asyncio.TaskGroup() as tg:
-            tg.create_task(sleep())
-            tg.create_task(work())
-            tg.create_task(explode())
-    except:
-        pass
-
-
-asyncio.run(main())
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)