tests/test_environment.py

import pytest
from server.environment import RedTeamEnvironment
from models import AttackAction, StrategyType, TargetCategory

@pytest.mark.asyncio
async def test_reset():
    env = RedTeamEnvironment()
    obs = await env.reset()
    assert obs.turn == 0
    assert obs.episode_id is not None
    assert env.is_active is True

@pytest.mark.asyncio
async def test_step():
    env = RedTeamEnvironment()
    await env.reset()
    action = AttackAction(
        strategy_type="roleplay",
        target_category="privacy",
        intensity=0.5,
        framing="Hello"
    )
    res = await env.step(action)
    assert res.observation.turn == 1
    assert len(env.attack_history) == 1

@pytest.mark.asyncio
async def test_guards():
    env = RedTeamEnvironment()
    action = AttackAction(strategy_type="roleplay", target_category="privacy", intensity=0.5, framing="H")
    with pytest.raises(ValueError, match="Episode not active"):
        await env.step(action)
    
    await env.reset()
    env.turn = 10
    with pytest.raises(ValueError, match="Episode already complete"):
        await env.step(action)

def test_framing_sanitization():
    action = AttackAction(
        strategy_type="roleplay", target_category="privacy", intensity=0.5,
        framing="<script>alert(1)</script> javascript:void(0)"
    )
    assert "<script>" not in action.framing
    assert "javascript:" not in action.framing
    assert "[FILTERED]" in action.framing

@pytest.mark.asyncio
async def test_injections():
    env = RedTeamEnvironment()
    
    called_reward = False
    def mock_reward(**kwargs):
        nonlocal called_reward
        called_reward = True
        return {"total_reward": 1.0}
    
    called_llm = False
    async def mock_llm(*args):
        nonlocal called_llm
        called_llm = True
        return {"defender_response": "ok", "attack_success": 0.5}
    
    env.set_reward_computer(mock_reward)
    env.set_llm_pipeline(mock_llm)
    
    await env.reset()
    action = AttackAction(strategy_type="roleplay", target_category="privacy", intensity=0.5, framing="H")
    await env.step(action)
    
    assert called_reward is True
    assert called_llm is True