85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock
|
|
import json
|
|
from aare.common.automation_queue_manager import WorkflowRedisManager
|
|
from aare.common.automation_models import (
|
|
QueueItem,
|
|
WorkflowEvent,
|
|
RuntimeState,
|
|
QueueItemStatus,
|
|
WorkflowStepRecord,
|
|
WorkflowStateKind,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_redis():
|
|
mock = MagicMock()
|
|
# mock pipeline
|
|
pipeline = MagicMock()
|
|
mock.pipeline.return_value = pipeline
|
|
pipeline.execute.return_value = []
|
|
return mock
|
|
|
|
@pytest.fixture
|
|
def manager(mock_redis):
|
|
return WorkflowRedisManager(mock_redis, beamline="x10sa")
|
|
|
|
def test_keys(manager):
|
|
assert manager._key("test") == "x10sa:workflow:test"
|
|
assert manager._item_key("123") == "x10sa:workflow:item:123"
|
|
assert manager._queue_key() == "x10sa:workflow:queue"
|
|
|
|
def test_create_item(manager, mock_redis):
|
|
item = QueueItem(item_id="item1", beamline="x10sa", owner_pgroup="p12345", samples=[], status=QueueItemStatus.PENDING, order_index=1, steps=[])
|
|
# Mock _next_item_id
|
|
mock_redis.incr.return_value = 1
|
|
mock_redis.zcard.return_value = 0
|
|
|
|
manager.create_item(item)
|
|
|
|
# Check if item was set in pipeline
|
|
pipeline = mock_redis.pipeline.return_value
|
|
pipeline.set.assert_called()
|
|
pipeline.zadd.assert_called()
|
|
|
|
def test_get_item(manager, mock_redis):
|
|
item = QueueItem(item_id="item1", beamline="x10sa", owner_pgroup="p12345", samples=[], status=QueueItemStatus.PENDING, order_index=1, steps=[])
|
|
mock_redis.get.return_value = item.model_dump_json()
|
|
|
|
ret_item = manager.get_item("item1")
|
|
assert ret_item.item_id == "item1"
|
|
mock_redis.get.assert_called_with("x10sa:workflow:item:item1")
|
|
|
|
def test_delete_item(manager, mock_redis):
|
|
manager.delete_item("item1")
|
|
pipeline = mock_redis.pipeline.return_value
|
|
pipeline.delete.assert_called_with("x10sa:workflow:item:item1")
|
|
pipeline.zrem.assert_called_with("x10sa:workflow:queue", "item1")
|
|
|
|
def test_get_runtime(manager, mock_redis):
|
|
runtime = RuntimeState(running=False)
|
|
mock_redis.get.return_value = runtime.model_dump_json()
|
|
|
|
ret_runtime = manager.get_runtime()
|
|
assert ret_runtime.running == False
|
|
|
|
def test_patch_runtime(manager, mock_redis):
|
|
runtime = RuntimeState(running=False)
|
|
mock_redis.get.return_value = runtime.model_dump_json()
|
|
|
|
manager.patch_runtime({"running": True})
|
|
mock_redis.set.assert_called()
|
|
# verify the content of set
|
|
args, kwargs = mock_redis.set.call_args
|
|
sent_data = json.loads(args[1])
|
|
assert sent_data["running"] == True
|
|
|
|
def test_append_event(manager, mock_redis):
|
|
event = WorkflowEvent(event_type="info", message="test event", beamline="x10sa")
|
|
manager.append_event(event)
|
|
mock_redis.xadd.assert_called()
|
|
# verify args
|
|
args, kwargs = mock_redis.xadd.call_args
|
|
assert args[0] == "x10sa:workflow:events"
|
|
assert "json" in args[1]
|