cleanup

Signed-off-by: Filinto Duran <[email protected]>

Author: filintod

durabletask/aio/ASYNCIO_ENHANCEMENTS.md

@@ -205,28 +205,17 @@ async def competitive_workflow(ctx, input_data):
 ```
 
 ### Error Handling with Context
+
 ```python
 async def robust_workflow(ctx, input_data):
     try:
         return await ctx.call_activity("risky_activity")
     except Exception as e:
         # Enhanced error will include workflow context
-        debug_info = ctx.get_debug_info()
+        debug_info = ctx._get_info_snapshot()
         return {"error": str(e), "debug": debug_info}
 ```
 
-### Cleanup Tasks
-```python
-async def workflow_with_cleanup(ctx, input_data):
-    async with ctx:  # Automatic cleanup
-        # Register cleanup tasks
-        ctx.add_cleanup(lambda: print("Workflow completed"))
-        
-        result = await ctx.call_activity("main_work")
-        return result
-    # Cleanup tasks run automatically here
-```
-
 ## Best Practices
 
 1. **Use deterministic alternatives**:
@@ -245,12 +234,7 @@ async def workflow_with_cleanup(ctx, input_data):
    result = await ctx.with_timeout(ctx.call_activity("external_api"), 30.0)
    ```
 
-4. **Use cleanup tasks for resource management**:
-   ```python
-   ctx.add_cleanup(lambda: cleanup_resources())
-   ```
-
-5. **Enable debug mode during development**:
+4. **Enable debug mode during development**:
    ```bash
    export DAPR_WF_DEBUG=true
    ```

durabletask/aio/context.py

@@ -337,15 +337,6 @@ async def __aexit__(
 
         self._cleanup_tasks.clear()
 
-    def add_cleanup(self, cleanup_fn: Callable[[], Any]) -> None:
-        """
-        Add a cleanup function to be called when the context exits.
-
-        Args:
-            cleanup_fn: Function to call during cleanup
-        """
-        self._cleanup_tasks.append(cleanup_fn)
-
     # Debug and monitoring
     def _log_operation(self, operation: str, details: Dict[str, Any]) -> None:
         """Log workflow operation for debugging."""
@@ -361,7 +352,7 @@ def _log_operation(self, operation: str, details: Dict[str, Any]) -> None:
             self._operation_history.append(entry)
             print(f"[WORKFLOW DEBUG] {operation}: {details}")
 
-    def get_debug_info(self) -> Dict[str, Any]:
+    def _get_info_snapshot(self) -> Dict[str, Any]:
         """
         Get debug information about the workflow execution.
 

tests/aio/test_asyncio_compat_enhanced.py renamed to tests/aio/test_asyncio_enhanced_additions.py

@@ -94,34 +94,12 @@ def test_get_debug_info(self):
             ctx = AsyncWorkflowContext(self.mock_base_ctx)
             ctx._log_operation("test_op", {"param": "value"})
 
-            debug_info = ctx.get_debug_info()
+            debug_info = ctx._get_info_snapshot()
 
             assert debug_info["instance_id"] == "test-instance-123"
             assert len(debug_info["operation_history"]) == 1
             assert debug_info["operation_history"][0]["type"] == "test_op"
 
-    def test_cleanup_registry(self):
-        cleanup_called = []
-
-        def cleanup_fn():
-            cleanup_called.append("sync")
-
-        async def async_cleanup_fn():
-            cleanup_called.append("async")
-
-        self.ctx.add_cleanup(cleanup_fn)
-        self.ctx.add_cleanup(async_cleanup_fn)
-
-        # Test cleanup execution
-        async def test_cleanup():
-            async with self.ctx:
-                pass
-
-        asyncio.run(test_cleanup())
-
-        # Cleanup should be called in reverse order
-        assert cleanup_called == ["async", "sync"]
-
     def test_activity_logging(self):
         with patch.dict(os.environ, {"DAPR_WF_DEBUG": "true"}):
             ctx = AsyncWorkflowContext(self.mock_base_ctx)
@@ -274,8 +252,6 @@ def test_strict_mode_secrets_blocking(self):
             pass
 
     def test_asyncio_sleep_patching(self):
-        import asyncio
-
         original_sleep = asyncio.sleep
 
         with _sandbox_scope(self.async_ctx, "best_effort"):

tests/aio/test_context.py

@@ -316,77 +316,14 @@ def test_operation_logging_in_debug_mode(self):
 
     def test_get_debug_info_method(self):
         """Test get_debug_info() method."""
-        debug_info = self.ctx.get_debug_info()
+        debug_info = self.ctx._get_info_snapshot()
 
         assert isinstance(debug_info, dict)
         assert debug_info["instance_id"] == "test-instance-123"
         assert debug_info["is_replaying"] == False
         assert "operation_history" in debug_info
         assert "cleanup_tasks_count" in debug_info
 
-    def test_add_cleanup_method(self):
-        """Test add_cleanup() method."""
-        cleanup_task = Mock()
-
-        self.ctx.add_cleanup(cleanup_task)
-
-        assert cleanup_task in self.ctx._cleanup_tasks
-
-    def test_async_context_manager(self):
-        """Test async context manager functionality."""
-        cleanup_task1 = Mock()
-        cleanup_task2 = Mock()
-
-        async def test_context_manager():
-            async with self.ctx:
-                self.ctx.add_cleanup(cleanup_task1)
-                self.ctx.add_cleanup(cleanup_task2)
-
-        # Run the async context manager
-        import asyncio
-
-        asyncio.run(test_context_manager())
-
-        # Cleanup tasks should have been called in reverse order
-        cleanup_task2.assert_called_once()
-        cleanup_task1.assert_called_once()
-
-    def test_async_context_manager_with_async_cleanup(self):
-        """Test async context manager with async cleanup tasks."""
-        import asyncio
-
-        async_cleanup = Mock()
-
-        async def _noop():
-            return None
-
-        async_cleanup.return_value = _noop()
-
-        async def test_async_cleanup():
-            async with self.ctx:
-                self.ctx.add_cleanup(async_cleanup)
-
-        # Should handle async cleanup tasks
-        asyncio.run(test_async_cleanup())
-
-    def test_async_context_manager_cleanup_error_handling(self):
-        """Test that cleanup errors don't prevent other cleanups."""
-        failing_cleanup = Mock(side_effect=Exception("Cleanup failed"))
-        working_cleanup = Mock()
-
-        async def test_cleanup_errors():
-            async with self.ctx:
-                self.ctx.add_cleanup(failing_cleanup)
-                self.ctx.add_cleanup(working_cleanup)
-
-        # Should not raise error and should call both cleanups
-        import asyncio
-
-        asyncio.run(test_cleanup_errors())
-
-        failing_cleanup.assert_called_once()
-        working_cleanup.assert_called_once()
-
     def test_detection_disabled_property(self):
         """Test _detection_disabled property."""
         import os

tests/aio/test_context_simple.py

@@ -15,7 +15,6 @@
 These tests focus on the actual implementation rather than expected features.
 """
 
-import asyncio
 import random
 import uuid
 from datetime import datetime, timedelta
@@ -261,34 +260,9 @@ def test_continue_as_new_method(self):
         # Should not raise error even if base context doesn't support it
         self.ctx.continue_as_new(new_input)
 
-    def test_add_cleanup_method(self):
-        """Test add_cleanup() method."""
-        cleanup_task = Mock()
-
-        self.ctx.add_cleanup(cleanup_task)
-
-        assert cleanup_task in self.ctx._cleanup_tasks
-
-    def test_async_context_manager(self):
-        """Test async context manager functionality."""
-        cleanup_task1 = Mock()
-        cleanup_task2 = Mock()
-
-        async def test_context_manager():
-            async with self.ctx:
-                self.ctx.add_cleanup(cleanup_task1)
-                self.ctx.add_cleanup(cleanup_task2)
-
-        # Run the async context manager
-        asyncio.run(test_context_manager())
-
-        # Cleanup tasks should have been called in reverse order
-        cleanup_task2.assert_called_once()
-        cleanup_task1.assert_called_once()
-
     def test_get_debug_info_method(self):
         """Test get_debug_info() method."""
-        debug_info = self.ctx.get_debug_info()
+        debug_info = self.ctx._get_info_snapshot()
 
         assert isinstance(debug_info, dict)
         assert debug_info["instance_id"] == "test-instance-123"