⚡️ Speed up function _temporal_context by 13% - #3

from __future__ import annotations

from contextvars import ContextVar
from typing import Optional, Union

# imports
import pytest  # used for our unit tests
from temporalio.nexus._operation_context import _temporal_context


# Dummy context classes for testing
class _TemporalStartOperationContext:
    def __init__(self, value):
        self.value = value

class _TemporalCancelOperationContext:
    def __init__(self, value):
        self.value = value

_temporal_start_operation_context: ContextVar[_TemporalStartOperationContext] = (
    ContextVar("temporal-start-operation-context")
)

_temporal_cancel_operation_context: ContextVar[_TemporalCancelOperationContext] = (
    ContextVar("temporal-cancel-operation-context")
)
from temporalio.nexus._operation_context import _temporal_context

# unit tests

# ----- BASIC TEST CASES -----





def test_both_contexts_set_raises():
    """Test that setting both contexts raises a RuntimeError."""
    start_ctx = _TemporalStartOperationContext("start")
    cancel_ctx = _TemporalCancelOperationContext("cancel")
    token1 = _temporal_start_operation_context.set(start_ctx)
    token2 = _temporal_cancel_operation_context.set(cancel_ctx)
    try:
        with pytest.raises(RuntimeError) as excinfo:
            _temporal_context()
    finally:
        _temporal_start_operation_context.reset(token1)
        _temporal_cancel_operation_context.reset(token2)

def test_neither_context_set_raises():
    """Test that not setting any context raises a RuntimeError."""
    with pytest.raises(RuntimeError) as excinfo:
        _temporal_context() # 625ns -> 542ns (15.3% faster)

def test_try_temporal_context_none():
    """Test that _try_temporal_context returns None if neither context is set."""

def test_context_reset_behavior():
    """Test that resetting contextvars properly removes context."""
    ctx = _TemporalStartOperationContext("start")
    token = _temporal_start_operation_context.set(ctx)
    _temporal_start_operation_context.reset(token)





def test_contextvars_isolation():
    """Test that contextvars are isolated between test runs (no leakage)."""
    # Set context, reset, and ensure next test sees nothing
    ctx = _TemporalStartOperationContext("isolate")
    token = _temporal_start_operation_context.set(ctx)
    _temporal_start_operation_context.reset(token)


#------------------------------------------------
from __future__ import annotations

from contextvars import ContextVar, Token
from typing import Optional, Union

# imports
import pytest  # used for our unit tests
from temporalio.nexus._operation_context import _temporal_context


# Dummy context classes for testing
class _TemporalStartOperationContext:
    def __init__(self, val: str = "start"):
        self.val = val

class _TemporalCancelOperationContext:
    def __init__(self, val: str = "cancel"):
        self.val = val

# ContextVars as used in the function
_temporal_start_operation_context: ContextVar[_TemporalStartOperationContext] = (
    ContextVar("temporal-start-operation-context")
)
_temporal_cancel_operation_context: ContextVar[_TemporalCancelOperationContext] = (
    ContextVar("temporal-cancel-operation-context")
)
from temporalio.nexus._operation_context import _temporal_context

# unit tests

# ----------- BASIC TEST CASES -----------




def test_temporal_context_none_set():
    """Test when neither context is set, should raise error."""
    with pytest.raises(RuntimeError, match="Not in Nexus operation context."):
        _temporal_context() # 792ns -> 667ns (18.7% faster)

# ----------- EDGE TEST CASES -----------


def test_temporal_context_contextvar_isolation():
    """Test that contextvars are isolated between tests (no residual state)."""



def test_temporal_context_contextvar_token_stack():
    """Test nested contextvar set/reset using tokens (contextvar stack)."""
    token1 = _temporal_start_operation_context.set(_TemporalStartOperationContext("first"))
    token2 = _temporal_start_operation_context.set(_TemporalStartOperationContext("second"))
    try:
        codeflash_output = _temporal_context(); ctx = codeflash_output
    finally:
        _temporal_start_operation_context.reset(token2)
        codeflash_output = _temporal_context(); ctx = codeflash_output
        _temporal_start_operation_context.reset(token1)

def test_temporal_context_contextvar_none_explicitly():
    """Test explicitly setting contextvar to None."""
    token = _temporal_start_operation_context.set(None)
    try:
        with pytest.raises(RuntimeError, match="Not in Nexus operation context."):
            _temporal_context()
    finally:
        _temporal_start_operation_context.reset(token)

# ----------- LARGE SCALE TEST CASES -----------




def test_temporal_context_large_scale_none_contexts():
    """Test that with no context set, always raises error, repeated."""
    for _ in range(100):
        with pytest.raises(RuntimeError, match="Not in Nexus operation context."):
            _temporal_context()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from temporalio.nexus._operation_context import _temporal_context
import pytest

def test__temporal_context():
    with pytest.raises(RuntimeError, match='Not\\ in\\ Nexus\\ operation\\ context\\.'):
        _temporal_context()