Skyvern/tests/unit/test_taskv2_block_context_preservation.py
Shuchang Zheng 76b10eb007
Fix OSS frontend build: add useFeatureFlag stub (#5042)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-10 23:36:42 -07:00

57 lines
2 KiB
Python

from skyvern.forge.sdk.core import skyvern_context
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
def test_loop_internal_state_dropped_by_naive_context_replacement():
"""Demonstrates the bug: creating a new context without carrying loop_internal_state drops it."""
original_context = SkyvernContext(
organization_id="org_1",
workflow_run_id="wr_1",
run_id="wr_1",
loop_internal_state={"downloaded_file_signatures_before_iteration": [("a.pdf", "abc", "https://files/a.pdf")]},
)
skyvern_context.set(original_context)
# Simulate TaskV2Block.finally creating a new context WITHOUT preserving loop state
context = skyvern_context.current()
new_context = SkyvernContext(
organization_id="org_1",
workflow_run_id="wr_1",
run_id=context.run_id,
)
skyvern_context.set(new_context)
result_context = skyvern_context.current()
# BUG: loop_internal_state is None
assert result_context.loop_internal_state is None
skyvern_context.reset()
def test_loop_internal_state_preserved_after_fix():
"""After fix: new context preserves loop_internal_state from old context."""
original_state = {"downloaded_file_signatures_before_iteration": [("a.pdf", "abc", "https://files/a.pdf")]}
original_context = SkyvernContext(
organization_id="org_1",
workflow_run_id="wr_1",
run_id="wr_1",
loop_internal_state=original_state,
)
skyvern_context.set(original_context)
# Simulate the FIXED TaskV2Block.finally
context = skyvern_context.current()
loop_state = context.loop_internal_state if context else None
new_context = SkyvernContext(
organization_id="org_1",
workflow_run_id="wr_1",
run_id=context.run_id,
loop_internal_state=loop_state,
)
skyvern_context.set(new_context)
result_context = skyvern_context.current()
assert result_context.loop_internal_state is not None
assert result_context.loop_internal_state == original_state
skyvern_context.reset()