"""
Tests for WorldModel.
"""

import pytest
from dcos.memory.world_model import WorldModel


class TestWorldModel:
    def test_add_and_get_entity(self):
        model = WorldModel()
        model.add_entity("e1", {"type": "robot", "status": "active"})
        entity = model.get_entity("e1")
        assert entity is not None
        assert entity["properties"]["type"] == "robot"

    def test_update_entity(self):
        model = WorldModel()
        model.add_entity("e1", {"status": "inactive"})
        model.update_entity("e1", {"status": "active"})
        entity = model.get_entity("e1")
        assert entity["properties"]["status"] == "active"

    def test_set_and_get_state(self):
        model = WorldModel()
        model.set_state("phase", "exploration")
        assert model.get_state("phase") == "exploration"

    def test_query_entities(self):
        model = WorldModel()
        model.add_entity("drone1", {"type": "drone", "team": "alpha"})
        model.add_entity("drone2", {"type": "drone", "team": "beta"})
        model.add_entity("robot1", {"type": "robot", "team": "alpha"})
        drones = model.query_entities("type", "drone")
        assert len(drones) == 2