3 files changed, 40 insertions, 24 deletions
diff --git a/t/random_wait.py b/t/manual/random_wait.py
index 2cfc290..2cfc290 100644
--- a/t/random_wait.py
+++ b/t/manual/random_wait.py
diff --git a/t/test_functionality.py b/t/test_functionality.py
index 12bf775..6daf2d3 100644
--- a/t/test_functionality.py
+++ b/t/test_functionality.py
@@ -1,6 +1,6 @@
 import utils
-import aiohttp
 import asyncio
+import threading
 
 
 class BasicUsage(utils.AergiaUnitTestCase):
@@ -54,7 +54,26 @@ class BasicUsage(utils.AergiaUnitTestCase):
 
         self.assertFuncContains('b', [self.expected_samples(delay * 3)],
                                 samples)
+        # TODO samples from gather all execution time, should we trace this??
+        self.assertFuncContains('a', [self.expected_samples(delay)], samples)
+
+    def test_alter_thread_task(self):
+        delay = 0.2
+
+        async def b(): await asyncio.sleep(delay)
+        async def a(): await asyncio.gather(b(), b(), b())
+        def c(): asyncio.run(a())
 
+        self.Aergia.start()
+        x = threading.Thread(target=c)
+        x.start()
+        x.join()
+        self.Aergia.stop()
+
+        samples = self.Aergia.get_samples()
+        self.assertFuncContains('c', [], samples)
+        self.assertFuncContains('b', [self.expected_samples(delay * 3)],
+                                samples)
         # TODO samples from gather all execution time, should we trace this??
         self.assertFuncContains('a', [self.expected_samples(delay)], samples)
 
@@ -71,3 +90,23 @@ class BasicUsage(utils.AergiaUnitTestCase):
 
         samples = self.Aergia.get_samples()
         self.assertFuncContains('a', [self.expected_samples(delay)], samples)
+
+    def test_async_gen_and_comp(self):
+        delay = 0.2
+        num_times = 10
+
+        async def b():
+            for i in range(num_times):
+                await asyncio.sleep(delay)
+                yield i
+
+        async def a():
+            return [r async for r in b()]
+
+        self.Aergia.start()
+        asyncio.run(a())
+        self.Aergia.stop()
+
+        samples = self.Aergia.get_samples()
+        self.assertFuncContains('b', [self.expected_samples(delay * num_times)],
+                                samples)
diff --git a/t/threads.py b/t/threads.py
deleted file mode 100644
index fa9c7ac..0000000
--- a/t/threads.py
+++ /dev/null
@@ -1,23 +0,0 @@
-import asyncio
-import threading
-
-
-
-async def count():
-    print("it's a secret!")
-    await asyncio.sleep(3)
-
-
-async def main():
-    await asyncio.gather(count(), count(), count())
-    print("done")
-
-
-def thread_func():
-    asyncio.run(main())
-
-
-if __name__ == "__main__":
-    x = threading.Thread(target=thread_func)
-    x.start()
-    x.join()