test_task_manager.py

# =====================================================
# Apckeyl Framework
# Version 1.0
# test_task_manager.py
# =====================================================

"""
Tests for Apckeyl Task Manager.

Version 1.0
"""

from task_manager import TaskManager


# =====================================================
# Test 1 — Task Manager Exists
# =====================================================

def test_task_manager_exists():

    manager = TaskManager()

    assert manager is not None


# =====================================================
# Test 2 — Create Image Upscale Task
# =====================================================

def test_create_image_upscale_task():

    manager = TaskManager()

    task = manager.create_task(
        task_type="image_upscale"
    )

    assert task is not None

    assert task["task_id"].startswith(
        "task_"
    )

    assert task["task_type"] == (
        "image_upscale"
    )

    assert task["status"] == (
        "created"
    )


# =====================================================
# Test 3 — Correct Module
# =====================================================

def test_task_module():

    manager = TaskManager()

    task = manager.create_task(
        task_type="image_upscale"
    )

    assert task["module_id"] == (
        "realesrgan"
    )

    assert task["module_name"] == (
        "Apckeyl_RealESRGAN"
    )


# =====================================================
# Test 4 — Payload
# =====================================================

def test_task_payload():

    manager = TaskManager()

    payload = {

        "filename": "test.png",

        "scale": 4,
    }

    task = manager.create_task(

        task_type="image_upscale",

        payload=payload,
    )

    assert task["payload"] == payload


# =====================================================
# Test 5 — Get Task
# =====================================================

def test_get_task():

    manager = TaskManager()

    task = manager.create_task(
        task_type="image_upscale"
    )

    loaded_task = manager.get_task(
        task["task_id"]
    )

    assert loaded_task is not None

    assert (
        loaded_task["task_id"]
        == task["task_id"]
    )


# =====================================================
# Test 6 — Update Status
# =====================================================

def test_update_status():

    manager = TaskManager()

    task = manager.create_task(
        task_type="image_upscale"
    )

    updated = manager.update_status(

        task["task_id"],

        "queued",
    )

    assert updated["status"] == (
        "queued"
    )


# =====================================================
# Test 7 — List Tasks
# =====================================================

def test_list_tasks():

    manager = TaskManager()

    manager.create_task(
        task_type="image_upscale"
    )

    manager.create_task(
        task_type="image_upscale"
    )

    tasks = manager.list_tasks()

    assert len(tasks) == 2


# =====================================================
# Test 8 — Unknown Task Rejected
# =====================================================

def test_unknown_task_rejected():

    manager = TaskManager()

    try:

        manager.create_task(
            task_type="unknown_task"
        )

    except LookupError:

        return

    raise AssertionError(
        "Unknown task was not rejected"
    )


# =====================================================
# Test 9 — Delete Task
# =====================================================

def test_delete_task():

    manager = TaskManager()

    task = manager.create_task(
        task_type="image_upscale"
    )

    manager.delete_task(
        task["task_id"]
    )

    assert (
        manager.get_task(
            task["task_id"]
        )
        is None
    )


# =====================================================
# Test Runner
# =====================================================

def main():

    tests = [

        test_task_manager_exists,

        test_create_image_upscale_task,

        test_task_module,

        test_task_payload,

        test_get_task,

        test_update_status,

        test_list_tasks,

        test_unknown_task_rejected,

        test_delete_task,
    ]

    passed = 0

    failed = 0

    print(
        "===================================="
    )

    print(
        "Apckeyl Task Manager Test"
    )

    print(
        "Version 1.0"
    )

    print(
        "===================================="
    )

    for test in tests:

        try:

            test()

            print(
                f"PASS: {test.__name__}"
            )

            passed += 1

        except Exception as error:

            print(
                f"FAIL: {test.__name__}"
            )

            print(
                f"ERROR: {error}"
            )

            failed += 1

    print(
        "===================================="
    )

    print(
        f"PASSED: {passed}"
    )

    print(
        f"FAILED: {failed}"
    )

    print(
        "===================================="
    )

    if failed > 0:

        raise RuntimeError(
            "Task Manager Test FAILED"
        )

    print(
        "TASK MANAGER TEST: PASSED"
    )


if __name__ == "__main__":

    main()