Reactor-Sim/tests/test_simulation.py

import json
from pathlib import Path

import pytest

from reactor_sim import constants
from reactor_sim.commands import ReactorCommand
from reactor_sim.failures import HealthMonitor
from reactor_sim.reactor import Reactor
from reactor_sim.simulation import ReactorSimulation


def test_reactor_initial_state_is_cold():
    reactor = Reactor.default()
    state = reactor.initial_state()
    assert state.core.fuel_temperature == constants.ENVIRONMENT_TEMPERATURE
    assert state.primary_loop.mass_flow_rate == 0.0
    assert state.total_electrical_output() == 0.0


def test_state_save_and_load_roundtrip(tmp_path: Path):
    reactor = Reactor.default()
    reactor.control.manual_control = True
    sim = ReactorSimulation(reactor, timestep=5.0, duration=15.0)
    sim.log()
    save_path = tmp_path / "plant_state.json"
    assert sim.last_state is not None
    reactor.save_state(str(save_path), sim.last_state)
    assert save_path.exists()
    restored_reactor = Reactor.default()
    restored_state = restored_reactor.load_state(str(save_path))
    assert restored_state.core.fuel_temperature == pytest.approx(
        sim.last_state.core.fuel_temperature
    )
    assert restored_reactor.control.rod_fraction == reactor.control.rod_fraction
    assert restored_reactor.control.manual_control == reactor.control.manual_control
    assert len(restored_state.primary_pumps) == 2
    assert len(restored_state.secondary_pumps) == 2


def test_health_monitor_flags_core_failure():
    reactor = Reactor.default()
    state = reactor.initial_state()
    state.core.fuel_temperature = constants.MAX_CORE_TEMPERATURE
    failures = reactor.health_monitor.evaluate(state, True, True, [True, True, True], dt=200.0)
    assert "core" in failures
    reactor._handle_failure("core")
    assert reactor.shutdown is True


def test_maintenance_recovers_component_health():
    monitor = HealthMonitor()
    pump = monitor.component("secondary_pump")
    pump.integrity = 0.3
    pump.fail()
    restored = monitor.maintain("secondary_pump", amount=0.5)
    assert restored is True
    assert pump.integrity == pytest.approx(0.8)
    assert pump.failed is False


def test_secondary_pump_loss_triggers_scram_and_no_steam():
    reactor = Reactor.default()
    state = reactor.initial_state()
    # Make sure some power is present to trigger heat-sink logic.
    state.core.power_output_mw = 500.0
    reactor.secondary_pump_active = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.1
    reactor.step(state, dt=1.0)
    assert reactor.shutdown is True
    assert all(t.electrical_output_mw == 0.0 for t in state.turbines)


def test_cold_shutdown_stays_subcritical():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.95
    reactor.primary_pump_active = False
    reactor.secondary_pump_active = False
    initial_power = state.core.power_output_mw
    for _ in range(10):
        reactor.step(state, dt=1.0)
    assert state.core.power_output_mw <= initial_power + 0.5
    assert reactor.shutdown is True


def test_toggle_maintenance_progresses_until_restored():
    reactor = Reactor.default()
    reactor.primary_pump_active = False
    pump = reactor.health_monitor.component("primary_pump")
    pump.integrity = 0.2

    def provider(t: float, _state):
        if t == 0:
            return ReactorCommand.maintain("primary_pump")
        return None

    sim = ReactorSimulation(reactor, timestep=1.0, duration=50.0, command_provider=provider)
    sim.log()
    assert pump.integrity >= 0.99
    assert "primary_pump" not in reactor.maintenance_active


def test_primary_pump_unit_toggle_updates_active_flag():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.primary_pump_active = True
    reactor.primary_pump_units = [True, True]

    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pumps={1: False}))
    assert reactor.primary_pump_units == [False, True]
    assert reactor.primary_pump_active is True

    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pumps={2: False}))
    assert reactor.primary_pump_units == [False, False]
    assert reactor.primary_pump_active is False


def test_secondary_pump_unit_toggle_can_restart_pump():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.secondary_pump_active = False
    reactor.secondary_pump_units = [False, False]

    reactor.step(state, dt=1.0, command=ReactorCommand(secondary_pumps={1: True}))

    assert reactor.secondary_pump_units == [True, False]
    assert reactor.secondary_pump_active is True
    assert state.secondary_loop.mass_flow_rate > 0.0


def test_primary_pumps_spool_up_over_seconds():
    reactor = Reactor.default()
    state = reactor.initial_state()
    # Enable both pumps and command full flow; spool should take multiple steps.
    target_flow = reactor.primary_pump.flow_rate(1.0) * len(reactor.primary_pump_units)
    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pump_on=True, coolant_demand=1.0))
    first_flow = state.primary_loop.mass_flow_rate
    assert 0.0 < first_flow < target_flow

    for _ in range(10):
        reactor.step(state, dt=1.0, command=ReactorCommand(coolant_demand=1.0))

    assert state.primary_loop.mass_flow_rate == pytest.approx(target_flow, rel=0.1)