Relicensed as Apache2.0

1 files changed, 40 insertions, 1 deletions

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)