Add a lot more tests, finish yappi comparisons

3 files changed, 209 insertions, 77 deletions

diff --git a/t/test_functionality.py b/t/test_functionality.py
index 2a2bbe3..eb9873a 100644
--- a/t/test_functionality.py
+++ b/t/test_functionality.py
@@ -41,11 +41,10 @@ class BasicUsage(utils.AergiaUnitTestCase):
                 tot = await b(tot, i)
             assert tot == 10
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
         yappi_samples = yappi.get_func_stats()
@@ -59,11 +58,10 @@ class BasicUsage(utils.AergiaUnitTestCase):
         async def b(): await asyncio.sleep(delay)
         async def a(): await asyncio.gather(b(), b(), b())
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
         yappi_samples = yappi.get_func_stats()
@@ -86,22 +84,30 @@ class BasicUsage(utils.AergiaUnitTestCase):
         async def a(): await asyncio.gather(b(), b(), b())
         def c(): asyncio.run(a())
 
+        yappi.start()
         self.Aergia.start()
         x = threading.Thread(target=c)
         x.start()
         x.join()
         self.Aergia.stop()
-        yappi.start()
-        x = threading.Thread(target=c)
-        x.start()
-        x.join()
         yappi.stop()
 
-        samples = self.Aergia.get_samples()
-        self.assert_reasonable_delay('c', 0, samples)
-        self.assert_reasonable_delay('b', delay * 3, samples)
+        yappi_samples = yappi.get_func_stats(ctx_id=0)
+        aergia_samples = self.Aergia.get_samples()
+
+        self.assert_reasonable_delay('c', 0, aergia_samples)
+        self.assert_reasonable_delay('b', delay * 3, aergia_samples)
         # see comment on `test_simultaneous_tasks'.
-        self.assert_reasonable_delay('a', delay, samples)
+        self.assert_reasonable_delay('a', delay, aergia_samples)
+
+        # Aergia does not assign time to the current task.
+        # Statistically, this means the select function, which traces
+        # up to the event loop (function c).
+        # Therefore this test would fail.
+        # self.assert_similar_delay('c', yappi_samples, aergia_samples)
+
+        self.assert_similar_delay('b', yappi_samples, aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
 
     def test_eager_task(self):
         delay = 0.2
@@ -110,15 +116,17 @@ class BasicUsage(utils.AergiaUnitTestCase):
             proc = await asyncio.create_subprocess_shell(f'sleep {delay}')
             await proc.communicate()
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
-        samples = self.Aergia.get_samples()
-        self.assert_reasonable_delay('a', delay, samples)
+        yappi_samples = yappi.get_func_stats(ctx_id=0)
+        aergia_samples = self.Aergia.get_samples()
+
+        self.assert_reasonable_delay('a', delay, aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
 
     def test_async_eager_scheduled_woven(self):
         d1 = 0.2
@@ -134,15 +142,17 @@ class BasicUsage(utils.AergiaUnitTestCase):
             proc = await asyncio.create_subprocess_shell(f'sleep {d4}')
             await proc.communicate()
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
-        samples = self.Aergia.get_samples()
-        self.assert_reasonable_delay('a', d1 + d2 + d3 + d4, samples)
+        yappi_samples = yappi.get_func_stats(ctx_id=0)
+        aergia_samples = self.Aergia.get_samples()
+
+        self.assert_reasonable_delay('a', d1 + d2 + d3 + d4, aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
 
     def test_async_generator(self):
         delay = 0.2
@@ -158,20 +168,25 @@ class BasicUsage(utils.AergiaUnitTestCase):
             async for item in b():
                 lst.append(item)
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
+        yappi.stop()
 
-        samples = self.Aergia.get_samples()
-        # TODO I do not think async generators report correctly.
-        # also does not work at all with yappi.
-        # doing so would be unique to this profiler.
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
 
-        # self.assert_reasonable_delay('b', delay * num_times, samples)
-        # self.assert_reasonable_delay('a', [], samples)
+        # TODO There are some minor shortcomings with the way Aergia handles
+        # async_generators. Lines which call the async generator will always
+        # be assigned the async generator's execution time as well. yappi
+        # also seems inconsistent.
+
+        self.assert_reasonable_delay('b', delay * num_times, aergia_samples)
+        self.assert_reasonable_delay('a', delay * num_times, aergia_samples)
 
-        self.assert_reasonable_delay('a', delay * num_times, samples)
-        self.assert_reasonable_delay('b', 0, samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
+        self.assert_similar_delay('b', yappi_samples, aergia_samples)
 
     def test_async_gen_and_comp(self):
         delay = 0.2
@@ -185,18 +200,24 @@ class BasicUsage(utils.AergiaUnitTestCase):
         async def a():
             return [r async for r in b()]
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
+        yappi.stop()
 
-        samples = self.Aergia.get_samples()
-        # TODO I do not think async generators report correctly.
-        # also does not work at all with yappi.
-        # doing so would be unique to this profiler.
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
 
-        self.assert_reasonable_delay('b', 0, samples)
-        self.assert_reasonable_delay('a', 0, samples)
-        self.assert_reasonable_delay('<listcomp>', delay * num_times, samples)
+        self.assert_reasonable_delay('b', delay * num_times, aergia_samples)
+        self.assert_reasonable_delay('a', 0, aergia_samples)
+        self.assert_reasonable_delay('<listcomp>', delay * num_times, aergia_samples)
+
+        self.assert_similar_delay('b', yappi_samples, aergia_samples)
+        # Aergia only assigns time to the current line when the task suspends
+        # This should fail.
+        # self.assert_similar_delay('a', yappi_samples, aergia_samples)
+        self.assert_similar_delay('<listcomp>', yappi_samples, aergia_samples)
 
     def test_deep_await(self):
         delay = 0.2
@@ -205,19 +226,24 @@ class BasicUsage(utils.AergiaUnitTestCase):
         async def b(): await c()
         async def a(): await b()
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
-        samples = self.Aergia.get_samples()
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
 
-        self.assert_reasonable_delay(
-            'c', delay, samples)
-        self.assert_reasonable_delay('b', 0, samples)
-        self.assert_reasonable_delay('a', 0, samples)
+        self.assert_reasonable_delay('c', delay, aergia_samples)
+        self.assert_reasonable_delay('b', 0, aergia_samples)
+        self.assert_reasonable_delay('a', 0, aergia_samples)
+
+        self.assert_similar_delay('c', yappi_samples, aergia_samples)
+        # Aergia does not assign time to the current task.
+        # These should fail.
+        # self.assert_similar_delay('b', yappi_samples, aergia_samples)
+        # self.assert_similar_delay('a', yappi_samples, aergia_samples)
 
     def test_task_groups(self):
         d1 = 0.2
@@ -233,25 +259,30 @@ class BasicUsage(utils.AergiaUnitTestCase):
                 tg.create_task(c())
                 tg.create_task(b())
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
-        samples = self.Aergia.get_samples()
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
 
-        self.assert_reasonable_delay('d', d1, samples)
-        self.assert_reasonable_delay('c', d2, samples)
-        self.assert_reasonable_delay('b', d3, samples)
+        self.assert_reasonable_delay('d', d1, aergia_samples)
+        self.assert_reasonable_delay('c', d2, aergia_samples)
+        self.assert_reasonable_delay('b', d3, aergia_samples)
         # the task group is technically waiting for all tasks to finish.
         # same as task.gather
         # This might be seen as unintuitive, (especially considering how
         # the next test works), though I don't want to bias the results
         # by adding logic to add artificial consistency.
         # Both are reporting correctly.
-        self.assert_reasonable_delay('a', d2, samples)
+        self.assert_reasonable_delay('a', d2, aergia_samples)
+
+        self.assert_similar_delay('d', yappi_samples, aergia_samples)
+        self.assert_similar_delay('c', yappi_samples, aergia_samples)
+        self.assert_similar_delay('b', yappi_samples, aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
 
     def test_task_groups_cancel(self):
         d1 = 0.1
@@ -273,37 +304,23 @@ class BasicUsage(utils.AergiaUnitTestCase):
             except* TerminateTaskGroup:
                 pass
 
+        yappi.start()
         self.Aergia.start()
         asyncio.run(a())
         self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a())
         yappi.stop()
 
-        samples = self.Aergia.get_samples()
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
 
-        self.assert_reasonable_delay('d', 0, samples)
-        self.assert_reasonable_delay('c', d1, samples)
-        self.assert_reasonable_delay('b', d3, samples)
+        self.assert_reasonable_delay('d', 0, aergia_samples)
+        self.assert_reasonable_delay('c', d1, aergia_samples)
+        self.assert_reasonable_delay('b', d3, aergia_samples)
         # this time is attached to the sleep call itself. Aergia's print
         # function would confirm this!
-        self.assert_reasonable_delay('a', d3, samples)
+        self.assert_reasonable_delay('a', d3, aergia_samples)
 
-    def test_asyncio_recursion(self):
-        delay = 0.1
-
-        async def a(n):
-            if n <= 0:
-                return
-            await asyncio.sleep(delay)
-            await a(n - 1)
-            await a(n - 2)
-
-        self.Aergia.start()
-        asyncio.run(a(3))
-        self.Aergia.stop()
-        yappi.start()
-        asyncio.run(a(3))
-        yappi.stop()
-
-        # samples = self.Aergia.get_samples()
+        self.assert_similar_delay('d', yappi_samples, aergia_samples)
+        self.assert_similar_delay('c', yappi_samples, aergia_samples)
+        self.assert_similar_delay('b', yappi_samples, aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
diff --git a/t/test_yappi_adaptations.py b/t/test_yappi_adaptations.py
new file mode 100644
index 0000000..7f034ef
--- /dev/null
+++ b/t/test_yappi_adaptations.py
@@ -0,0 +1,112 @@
+import yappi
+import utils
+import asyncio
+import threading
+
+# A test file containing tests adapted from yappi's own test suite.
+
+
+class YappiTests(utils.AergiaUnitTestCase):
+
+    def test_asyncio_recursion_yappi(self):
+        delay = 0.1
+
+        async def a(n):
+            if n <= 0:
+                return
+            await asyncio.sleep(delay)
+            await a(n - 1)
+            await a(n - 2)
+
+        yappi.start()
+        self.Aergia.start()
+        asyncio.run(a(3))
+        self.Aergia.stop()
+        yappi.stop()
+
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
+
+        self.assert_reasonable_delay('a', delay * 4, aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
+
+    def test_basic_multithread(self):
+        delay = 0.1
+        num_times = 5
+
+        async def a():
+            await asyncio.sleep(delay)
+
+        async def b():
+            await a()
+
+        async def recursive_a(n):
+            if not n:
+                return
+            await asyncio.sleep(delay)
+            await recursive_a(n - 1)
+
+        def tag_cbk():
+            cthread = threading.current_thread()
+            try:
+                return cthread._tag
+            except:
+                return -1
+
+        threading.current_thread()._tag = 0
+        yappi.set_tag_callback(tag_cbk)
+
+        def _thread_event_loop(loop):
+            asyncio.set_event_loop(loop)
+            loop.run_forever()
+
+        _TCOUNT = 3
+        _ctag = 1
+
+        ts = []
+        for i in range(_TCOUNT):
+            _loop = asyncio.new_event_loop()
+            t = threading.Thread(target=_thread_event_loop, args=(_loop, ))
+            t._tag = _ctag
+            t._loop = _loop
+            t.start()
+
+            ts.append(t)
+            _ctag += 1
+
+        async def stop_loop():
+            asyncio.get_event_loop().stop()
+
+        async def driver():
+            futs = []
+            fut = asyncio.run_coroutine_threadsafe(a(), ts[0]._loop)
+            futs.append(fut)
+            fut = asyncio.run_coroutine_threadsafe(recursive_a(num_times), ts[1]._loop)
+            futs.append(fut)
+            fut = asyncio.run_coroutine_threadsafe(b(), ts[2]._loop)
+            futs.append(fut)
+            for fut in futs:
+                fut.result()
+
+            for t in ts:
+                asyncio.run_coroutine_threadsafe(stop_loop(), t._loop)
+
+        yappi.start()
+        self.Aergia.start()
+        asyncio.run(driver())
+        self.Aergia.stop()
+        yappi.stop()
+
+        yappi_samples = yappi.get_func_stats()
+        aergia_samples = self.Aergia.get_samples()
+
+        self.assert_reasonable_delay('a', delay * 2, aergia_samples)
+        self.assert_reasonable_delay('b', 0, aergia_samples)
+        self.assert_reasonable_delay('recursive_a',
+                                     delay * num_times,
+                                     aergia_samples)
+        self.assert_similar_delay('a', yappi_samples, aergia_samples)
+        self.assert_similar_delay('recursive_a', yappi_samples, aergia_samples)
+        # Aergia only assigns time to the current line when the task suspends
+        # This should fail.
+        # self.assert_similar_delay('b', yappi_samples, aergia_samples)
diff --git a/t/utils.py b/t/utils.py
index 9d77665..4dde466 100644
--- a/t/utils.py
+++ b/t/utils.py
@@ -27,6 +27,8 @@ class AergiaUnitTestCase(unittest.TestCase):
         self.assert_roughly_equal(time_actual, time_expected)
 
     def assert_similar_delay(self, func_name, yappi_stats, aergia_samples):
+        '''Compares the results reported by Aergia for FUNC_NAME
+        with yappi.'''
         time_yappi = self.yappi_extract_values_by_func(yappi_stats, func_name)
 
         samples_aergia = self.aergia_extract_values_by_func(
@@ -36,10 +38,11 @@ class AergiaUnitTestCase(unittest.TestCase):
         self.assert_roughly_equal(time_aergia, time_yappi)
 
     def assert_roughly_equal(self, v1, v2):
-        '''Throws an exception if values V1 and V2 are not within 0.03
-        seconds of each other.'''
+        '''Throws an exception if values V1 and V2 are not within 0.035
+        seconds of each other. This number is influenced by instrumentation
+        overhead, sampling inconsistency, etc.'''
         a = abs(v1 - v2)
-        self.assertTrue(a <= .03, f'{v1} (expected) not roughly {v2} (actual)')
+        self.assertTrue(a <= .035, f'{v1} (actual) not roughly {v2} (expected)')
 
     def aergia_expected_time(self, total_samples):
         '''Given TOTAL_SAMPLES, returns the total time, using the