from radiuma_api.scheduler import Scheduler
from radiuma_api.workflow import Workflow
from radiuma_api.module import Module
from radiuma_api.io_port import InPort, OutPort, ImageType
from radiuma_api.execution_context import ExecutionContext
from radiuma_api.task import Status
# Codes Go below:


class Reader(Module):
    def __init__(self):
        super().__init__("Reader")
        self._add_output_port(OutPort("images", ImageType()))

    def run(self, context):
        context.put(self.outputs[0], "DATA")
        self._set_status(Status.COMPLETED)

class Extractor(Module):
    def __init__(self):
        super().__init__("Extractor")
        self._add_input_port(InPort("images", ImageType()))

    def run(self, context):
        assert context.get(self.inputs[0]) == "DATA"
        self._set_status(Status.COMPLETED)

def test_scheduler_runs_in_order():
    ctx = ExecutionContext()
    wf = Workflow("WF")

    r = Reader()
    e = Extractor()

    r.outputs[0].connect(e.inputs[0])

    wf.add_child(r)
    wf.add_child(e)

    wf.run(ctx)

    assert r.status == Status.COMPLETED
    assert e.status == Status.COMPLETED