mirror of
https://github.com/trailofbits/buttercup
synced 2026-06-21 14:11:39 +00:00
Riccardo Schirone
47f38b8863
Run `ruff check --fix` across all components to fix I001 import sorting violations in src/ and test/ directories. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
from buttercup.common.datastructures.msg_pb2 import IndexRequest
|
|
from buttercup.common.task_registry import TaskRegistry
|
|
from redis import Redis
|
|
|
|
from buttercup.program_model.program_model import ProgramModel
|
|
|
|
|
|
@pytest.fixture
|
|
def redis_client():
|
|
res = Mock(spec=Redis)
|
|
return res
|
|
|
|
|
|
@pytest.fixture
|
|
def program_model(redis_client):
|
|
model = ProgramModel(redis=redis_client)
|
|
# Ensure the queues are mocked since we won't have a real Redis
|
|
model.task_queue = Mock()
|
|
model.output_queue = Mock()
|
|
model.registry = Mock(spec=TaskRegistry)
|
|
return model
|
|
|
|
|
|
def test_serve_item_skip_cancelled_task(program_model):
|
|
"""Test that cancelled or expired tasks are skipped"""
|
|
# Set up the mocks
|
|
mock_task_id = "test_cancelled_task_id"
|
|
|
|
# Create mock IndexRequest
|
|
mock_request = IndexRequest(task_id=mock_task_id)
|
|
|
|
# Create mock RQItem
|
|
mock_item = MagicMock()
|
|
mock_item.deserialized = mock_request
|
|
mock_item.item_id = "test_item_id"
|
|
|
|
# Mock the queue to return our test item
|
|
program_model.task_queue.pop.return_value = mock_item
|
|
|
|
# Mock registry to indicate the task should be stopped (cancelled or expired)
|
|
program_model.registry.should_stop_processing.return_value = True
|
|
|
|
# Call serve_item
|
|
result = program_model.serve_item()
|
|
|
|
# Verify the task was acknowledged without processing
|
|
assert result is True
|
|
program_model.registry.should_stop_processing.assert_called_once_with(mock_task_id)
|
|
program_model.task_queue.ack_item.assert_called_once_with(mock_item.item_id)
|
|
|
|
# Verify process_task was not called
|
|
# This is a bit tricky since it's a method on the same object
|
|
# We'll patch it for the next test to verify it's called for normal tasks
|
|
|
|
|
|
def test_serve_item_process_normal_task(program_model):
|
|
"""Test that normal tasks are processed"""
|
|
# Set up the mocks
|
|
mock_task_id = "test_normal_task_id"
|
|
|
|
# Create mock IndexRequest
|
|
mock_request = IndexRequest(task_id=mock_task_id)
|
|
|
|
# Create mock RQItem
|
|
mock_item = MagicMock()
|
|
mock_item.deserialized = mock_request
|
|
mock_item.item_id = "test_item_id"
|
|
|
|
# Mock the queue to return our test item
|
|
program_model.task_queue.pop.return_value = mock_item
|
|
|
|
# Mock registry to indicate the task should not be stopped
|
|
program_model.registry.should_stop_processing.return_value = False
|
|
|
|
# Mock process_task to return success
|
|
with patch.object(program_model, "process_task", return_value=True):
|
|
# Call serve_item
|
|
result = program_model.serve_item()
|
|
|
|
# Verify the task was processed and acknowledged
|
|
assert result is True
|
|
program_model.registry.should_stop_processing.assert_called_once_with(mock_task_id)
|
|
program_model.process_task.assert_called_once_with(mock_request)
|
|
program_model.task_queue.ack_item.assert_called_once_with(mock_item.item_id)
|
|
program_model.output_queue.push.assert_called_once()
|