"""Tests for InMemoryStateBackend."""

from fusionagi.core.memory_backend import InMemoryStateBackend
from fusionagi.schemas.task import Task, TaskState


def test_set_and_get():
    backend = InMemoryStateBackend()
    task = Task(task_id="m1", goal="memory test")
    backend.set_task(task)
    assert backend.get_task("m1") is not None
    assert backend.get_task("m1").goal == "memory test"


def test_state_management():
    backend = InMemoryStateBackend()
    backend.set_task(Task(task_id="m2", goal="state"))
    backend.set_task_state("m2", TaskState.ACTIVE)
    assert backend.get_task_state("m2") == TaskState.ACTIVE


def test_traces():
    backend = InMemoryStateBackend()
    backend.set_task(Task(task_id="m3", goal="traces"))
    backend.append_trace("m3", {"a": 1})
    backend.append_trace("m3", {"b": 2})
    assert len(backend.get_trace("m3")) == 2


def test_delete():
    backend = InMemoryStateBackend()
    backend.set_task(Task(task_id="m4", goal="del"))
    assert backend.delete_task("m4") is True
    assert backend.delete_task("m4") is False


def test_list_and_count():
    backend = InMemoryStateBackend()
    for i in range(3):
        backend.set_task(Task(task_id=f"l{i}", goal=f"g{i}"))
    assert backend.count_tasks() == 3
    assert len(backend.list_tasks()) == 3