"""
Tests for the TaskScheduler.
"""

import pytest
from dcos.core.scheduler import TaskScheduler, Task


class TestTaskScheduler:
    def test_enqueue_dequeue(self):
        scheduler = TaskScheduler()
        task = Task(name="test", agent_id="agent_1", priority=5)
        scheduler.enqueue(task)
        dequeued = scheduler.dequeue()
        assert dequeued is not None
        assert dequeued.name == "test"
        assert dequeued.status.name == "RUNNING"

    def test_priority_ordering(self):
        scheduler = TaskScheduler()
        t1 = Task(name="low", agent_id="a1", priority=1)
        t2 = Task(name="high", agent_id="a2", priority=10)
        scheduler.enqueue(t1)
        scheduler.enqueue(t2)
        first = scheduler.dequeue()
        assert first is not None and first.name == "high"

    def test_complete_task(self):
        scheduler = TaskScheduler()
        task = Task(name="done", agent_id="a1")
        scheduler.enqueue(task)
        dequeued = scheduler.dequeue()
        assert dequeued is not None
        result = scheduler.complete(dequeued.id)
        assert result is True

    def test_dependency_wait(self):
        scheduler = TaskScheduler()
        dep = Task(name="dependency", agent_id="a1")
        main = Task(name="main", agent_id="a2", depends_on=[dep.id])
        scheduler.enqueue(dep)
        scheduler.enqueue(main)
        # Only dep should dequeue since main depends on dep
        first = scheduler.dequeue()
        assert first is not None and first.name == "dependency"
        second = scheduler.dequeue()
        assert second is None  # main still blocked

    def test_empty_queue(self):
        scheduler = TaskScheduler()
        assert scheduler.dequeue() is None