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?

3 files changed, 69 insertions, 6 deletions

diff --git a/t/async-generator-and-comprehension.py b/t/async-generator-and-comprehension.py
index ce85f45..2e1b60b 100644
--- a/t/async-generator-and-comprehension.py
+++ b/t/async-generator-and-comprehension.py
@@ -1,4 +1,5 @@
 import asyncio
+import yappi
 
 
 async def async_generator():
@@ -12,4 +13,8 @@ async def main():
     print(sum(r))
 
 
+yappi.set_clock_type("wall")
+yappi.start()
 asyncio.run(main())
+yappi.get_func_stats().print_all()
+yappi.get_thread_stats().print_all()
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)
diff --git a/t/utils.py b/t/utils.py
index 5772fac..da9a9c1 100644
--- a/t/utils.py
+++ b/t/utils.py
@@ -23,7 +23,7 @@ class AergiaUnitTestCase(unittest.TestCase):
 
     def assertRoughlyEqual(self, v1, v2):
         a = abs(v1 - v2)
-        self.assertTrue(a <= 1, f'{v1} (expected) not roughly {v2} (actual)')
+        self.assertTrue(a <= 2, f'{v1} (expected) not roughly {v2} (actual)')
 
     def expected_samples(self, total_seconds):
         return (total_seconds / self.interval)