mirror of
https://github.com/trailofbits/buttercup
synced 2026-06-21 14:11:39 +00:00
Dan Guido
42691e50b4
* style: apply ruff auto-fixes and formatting across entire codebase Applied safe auto-fixes from ruff v0.12.9 with --select ALL to improve code quality: - Reorder imports (stdlib → third-party → local) - Use modern type hints (collections.abc.Generator instead of typing.Generator) - Add trailing commas for better diffs - Format multi-line function parameters for readability - Add strict=False to zip() calls for explicit behavior - Simplify redundant elif to if after return statements - Consistent code formatting with ruff format These are all mechanical, non-controversial changes that improve code consistency without altering functionality. Changes affect 180 files across all modules: common, fuzzer, orchestrator, patcher, program-model, and seed-gen. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com> * re-applt ruff after merge --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Michael D Brown <michael.brown@trailofbits.com>
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
from redis import Redis
|
|
|
|
from buttercup.common.datastructures.msg_pb2 import IndexRequest
|
|
from buttercup.common.task_registry import TaskRegistry
|
|
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()
|