Aergia.get_deepest_traceable_frame refactor, tests

Generators and list comprehension gives correct but less ideal results, matching yappi. Can this be improved in the future?

1 files changed, 63 insertions, 5 deletions

diff --git a/t/test_functionality.py b/t/test_functionality.py
index 6daf2d3..b60cb46 100644
--- a/t/test_functionality.py
+++ b/t/test_functionality.py
@@ -38,7 +38,8 @@ class BasicUsage(utils.AergiaUnitTestCase):
 
         samples = self.Aergia.get_samples()
 
-        self.assertFuncContains('b', [self.expected_samples(delay * num_times)],
+        self.assertFuncContains('b',
+                                [self.expected_samples(delay * num_times)],
                                 samples)
 
     def test_simultaneous_tasks(self):
@@ -54,10 +55,10 @@ class BasicUsage(utils.AergiaUnitTestCase):
 
         self.assertFuncContains('b', [self.expected_samples(delay * 3)],
                                 samples)
-        # TODO samples from gather all execution time, should we trace this??
+        # 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):
+    def test_subthread_task(self):
         delay = 0.2
 
         async def b(): await asyncio.sleep(delay)
@@ -74,7 +75,7 @@ class BasicUsage(utils.AergiaUnitTestCase):
         self.assertFuncContains('c', [], samples)
         self.assertFuncContains('b', [self.expected_samples(delay * 3)],
                                 samples)
-        # TODO samples from gather all execution time, should we trace this??
+        # TODO samples from gather all execution time, should we trace this?
         self.assertFuncContains('a', [self.expected_samples(delay)], samples)
 
     def test_eager_task(self):
@@ -91,6 +92,39 @@ class BasicUsage(utils.AergiaUnitTestCase):
         samples = self.Aergia.get_samples()
         self.assertFuncContains('a', [self.expected_samples(delay)], samples)
 
+    def test_async_generator(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():
+            lst = []
+            async for item in b():
+                lst.append(item)
+
+        self.Aergia.start()
+        asyncio.run(a())
+        self.Aergia.stop()
+
+        samples = self.Aergia.get_samples()
+        # TODO can we make these results more intuitive?
+        # async generators also do not work at all with yappi.
+        # doing so would be unique to this profiler.
+
+        # self.assertFuncContains('b',
+        #                         [self.expected_samples(delay * num_times)],
+        #                         samples)
+        # self.assertFuncContains('a', [], samples)
+
+        self.assertFuncContains('a',
+                                [self.expected_samples(delay * num_times)],
+                                samples)
+        self.assertFuncContains('b', [], samples)
+
     def test_async_gen_and_comp(self):
         delay = 0.2
         num_times = 10
@@ -108,5 +142,29 @@ class BasicUsage(utils.AergiaUnitTestCase):
         self.Aergia.stop()
 
         samples = self.Aergia.get_samples()
-        self.assertFuncContains('b', [self.expected_samples(delay * num_times)],
+        # TODO can we make these results more intuitive?
+        # async generators also do not work at all with yappi.
+        # doing so would be unique to this profiler.
+
+        self.assertFuncContains('b', [], samples)
+        self.assertFuncContains('a', [], samples)
+        self.assertFuncContains('<listcomp>',
+                                [self.expected_samples(delay * num_times)],
                                 samples)
+
+    def test_deep_await(self):
+        delay = 0.2
+
+        async def c(): await asyncio.sleep(delay)
+        async def b(): await c()
+        async def a(): await b()
+
+        self.Aergia.start()
+        asyncio.run(a())
+        self.Aergia.stop()
+
+        samples = self.Aergia.get_samples()
+
+        self.assertFuncContains('c', [self.expected_samples(delay)], samples)
+        self.assertFuncContains('b', [], samples)
+        self.assertFuncContains('a', [], samples)