PrimAITE/tests/test_resetting_node.py

# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
"""Used to test Active Node functions."""
import pytest

from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState
from primaite.common.service import Service
from primaite.config.training_config import TrainingConfig
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode


@pytest.mark.parametrize(
    "starting_operating_state, expected_operating_state",
    [(HardwareState.RESETTING, HardwareState.ON)],
)
def test_node_resets_correctly(starting_operating_state, expected_operating_state):
    """Tests that a node resets correctly."""
    active_node = ActiveNode(
        node_id="0",
        name="node",
        node_type=NodeType.COMPUTER,
        priority=Priority.P1,
        hardware_state=starting_operating_state,
        ip_address="192.168.0.1",
        software_state=SoftwareState.COMPROMISED,
        file_system_state=FileSystemState.CORRUPT,
        config_values=TrainingConfig(),
    )

    for x in range(5):
        active_node.update_resetting_status()

    assert active_node.software_state == SoftwareState.GOOD
    assert active_node.file_system_state_actual == FileSystemState.GOOD
    assert active_node.hardware_state == expected_operating_state


@pytest.mark.parametrize(
    "operating_state, expected_operating_state",
    [(HardwareState.BOOTING, HardwareState.ON)],
)
def test_node_boots_correctly(operating_state, expected_operating_state):
    """Tests that a node boots correctly."""
    service_node = ServiceNode(
        node_id=0,
        name="node",
        node_type="COMPUTER",
        priority="1",
        hardware_state=operating_state,
        ip_address="192.168.0.1",
        software_state=SoftwareState.GOOD,
        file_system_state="GOOD",
        config_values=1,
    )
    service_attributes = Service(name="node", port="80", software_state=SoftwareState.COMPROMISED)
    service_node.add_service(service_attributes)

    for x in range(5):
        service_node.update_booting_status()

    assert service_attributes.software_state == SoftwareState.GOOD
    assert service_node.hardware_state == expected_operating_state


@pytest.mark.parametrize(
    "operating_state, expected_operating_state",
    [(HardwareState.SHUTTING_DOWN, HardwareState.OFF)],
)
def test_node_shutdown_correctly(operating_state, expected_operating_state):
    """Tests that a node shutdown correctly."""
    active_node = ActiveNode(
        node_id=0,
        name="node",
        node_type="COMPUTER",
        priority="1",
        hardware_state=operating_state,
        ip_address="192.168.0.1",
        software_state=SoftwareState.GOOD,
        file_system_state="GOOD",
        config_values=1,
    )

    for x in range(5):
        active_node.update_shutdown_status()

    assert active_node.hardware_state == expected_operating_state