From e60918324233e99c30763484f0da2ac6673f2f89 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Mon, 29 Aug 2022 15:08:04 -0400
Subject: [PATCH] Expose lifetime stack as class attr, add base test suite

---
 tests/test_runtime.py  | 70 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_spawning.py |  1 +
 tractor/_runtime.py    | 12 ++++----
 3 files changed, 77 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_runtime.py

diff --git a/tests/test_runtime.py b/tests/test_runtime.py
new file mode 100644
index 0000000..5604d6f
--- /dev/null
+++ b/tests/test_runtime.py
@@ -0,0 +1,70 @@
+"""
+Verifying internal runtime state and undocumented extras.
+
+"""
+import os
+
+import pytest
+import trio
+import tractor
+
+from conftest import tractor_test
+
+
+_file_path: str = ''
+
+
+def unlink_file():
+    print('Removing tmp file!')
+    os.remove(_file_path)
+
+
+async def crash_and_clean_tmpdir(
+    tmp_file_path: str,
+    error: bool = True,
+):
+    global _file_path
+    _file_path = tmp_file_path
+
+    actor = tractor.current_actor()
+    actor.lifetime_stack.callback(unlink_file)
+
+    assert os.path.isfile(tmp_file_path)
+    await trio.sleep(0.1)
+    if error:
+        assert 0
+    else:
+        actor.cancel_soon()
+
+
+@pytest.mark.parametrize(
+    'error_in_child',
+    [True, False],
+)
+@tractor_test
+async def test_lifetime_stack_wipes_tmpfile(
+    tmp_path,
+    error_in_child: bool,
+):
+    child_tmp_file = tmp_path / "child.txt"
+    child_tmp_file.touch()
+    assert child_tmp_file.exists()
+    path = str(child_tmp_file)
+
+    try:
+        with trio.move_on_after(0.5):
+            async with tractor.open_nursery() as n:
+                    await (  # inlined portal
+                        await n.run_in_actor(
+                            crash_and_clean_tmpdir,
+                            tmp_file_path=path,
+                            error=error_in_child,
+                        )
+                    ).result()
+
+    except tractor.RemoteActorError:
+        pass
+
+    # tmp file should have been wiped by
+    # teardown stack.
+    assert not child_tmp_file.exists()
diff --git a/tests/test_spawning.py b/tests/test_spawning.py
index f1679f6..3e5ab1b 100644
--- a/tests/test_spawning.py
+++ b/tests/test_spawning.py
@@ -1,5 +1,6 @@
 """
 Spawning basics
+
 """
 from typing import Optional
 
diff --git a/tractor/_runtime.py b/tractor/_runtime.py
index 4634375..c2f9551 100644
--- a/tractor/_runtime.py
+++ b/tractor/_runtime.py
@@ -337,10 +337,6 @@ def _get_mod_abspath(module):
     return os.path.abspath(module.__file__)
 
 
-# process-global stack closed at end on actor runtime teardown
-_lifetime_stack: ExitStack = ExitStack()
-
-
 async def try_ship_error_to_parent(
     channel: Channel,
     err: Union[Exception, trio.MultiError],
@@ -406,6 +402,10 @@ class Actor:
     # if started on ``asycio`` running ``trio`` in guest mode
     _infected_aio: bool = False
 
+    # Process-global stack closed at end on actor runtime teardown.
+    # NOTE: this is currently an undocumented public api.
+    lifetime_stack: ExitStack = ExitStack()
+
     def __init__(
         self,
         name: str,
@@ -1293,7 +1293,7 @@ async def async_main(
         # killed (i.e. this actor is cancelled or signalled by the parent)
     except Exception as err:
         log.info("Closing all actor lifetime contexts")
-        _lifetime_stack.close()
+        actor.lifetime_stack.close()
 
         if not registered_with_arbiter:
             # TODO: I guess we could try to connect back
@@ -1332,7 +1332,7 @@ async def async_main(
         #     with trio.CancelScope(shield=True):
         #         await _debug.breakpoint()
 
-        _lifetime_stack.close()
+        actor.lifetime_stack.close()
 
         # Unregister actor from the arbiter
         if registered_with_arbiter and (