admin/trailofbits-buttercup
1
0
Fork 0
mirror of https://github.com/trailofbits/buttercup synced 2026-06-21 14:11:39 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
pre-commit-infrastructure
trailofbits-buttercup/fuzzer/tests/test_builder_bot.py
T
Dan Guido 42691e50b4 style: apply ruff auto-fixes and formatting across entire codebase (#309) 
* 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>
2025-08-22 10:30:12 -04:00

589 lines
24 KiB
Python
Raw Permalink Blame History

import unittest
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildRequest, BuildType
from buttercup.common.task_registry import TaskRegistry
from buttercup.fuzzing_infra.builder_bot import BuilderBot
class TestBuilderBot(unittest.TestCase):
def setUp(self):
self.redis_mock = Mock(spec=Redis)
self.seconds_sleep = 1.0
self.allow_caching = True
self.allow_pull = True
self.python = "/usr/bin/python3"
self.wdir = "/tmp/test_wdir"
# Mock queue factory and queues
self.queue_factory_mock = Mock()
self.build_requests_queue_mock = Mock()
self.build_outputs_queue_mock = Mock()
self.queue_factory_mock.create.side_effect = [self.build_requests_queue_mock, self.build_outputs_queue_mock]
self.task_registry_mock = MagicMock(spec=TaskRegistry)
self.task_registry_mock.should_stop_processing.return_value = False
# Create BuilderBot instance with mocked dependencies
with (
patch("buttercup.fuzzing_infra.builder_bot.QueueFactory") as queue_factory_class_mock,
patch("buttercup.fuzzing_infra.builder_bot.TaskRegistry") as task_registry_class_mock,
):
queue_factory_class_mock.return_value = self.queue_factory_mock
task_registry_class_mock.return_value = self.task_registry_mock
self.builder_bot = BuilderBot(
redis=self.redis_mock,
seconds_sleep=self.seconds_sleep,
allow_caching=self.allow_caching,
allow_pull=self.allow_pull,
python=self.python,
wdir=self.wdir,
)
def test_post_init_sets_up_queues(self):
"""Test that __post_init__ correctly sets up the queues."""
# Verify that the queues were created with correct parameters
self.assertEqual(self.builder_bot._build_requests_queue, self.build_requests_queue_mock)
self.assertEqual(self.builder_bot._build_outputs_queue, self.build_outputs_queue_mock)
self.assertEqual(self.queue_factory_mock.create.call_count, 2)
def test_apply_challenge_diff_no_diff(self):
"""Test _apply_challenge_diff when no diff is requested."""
task_mock = Mock()
msg = BuildRequest(apply_diff=False)
result = self.builder_bot._apply_challenge_diff(task_mock, msg)
self.assertTrue(result)
task_mock.apply_patch_diff.assert_not_called()
def test_apply_challenge_diff_with_diff_success(self):
"""Test _apply_challenge_diff when diff is requested and succeeds."""
task_mock = Mock()
task_mock.apply_patch_diff.return_value = True
msg = BuildRequest(apply_diff=True)
result = self.builder_bot._apply_challenge_diff(task_mock, msg)
self.assertTrue(result)
task_mock.apply_patch_diff.assert_called_once()
def test_apply_challenge_diff_with_diff_failure(self):
"""Test _apply_challenge_diff when diff is requested but fails."""
task_mock = Mock()
task_mock.apply_patch_diff.return_value = False
msg = BuildRequest(apply_diff=True)
result = self.builder_bot._apply_challenge_diff(task_mock, msg)
self.assertFalse(result)
task_mock.apply_patch_diff.assert_called_once()
def test_apply_patch_no_patch(self):
"""Test _apply_patch when no patch is provided."""
task_mock = Mock()
msg = BuildRequest(patch="")
result = self.builder_bot._apply_patch(task_mock, msg)
self.assertTrue(result)
task_mock.apply_patch_diff.assert_not_called()
@patch("tempfile.NamedTemporaryFile")
def test_apply_patch_with_patch_success(self, temp_file_mock):
"""Test _apply_patch when patch is provided and succeeds."""
task_mock = Mock()
task_mock.apply_patch_diff.return_value = True
patch_content = "--- a/file.c\n+++ b/file.c\n@@ -1,3 +1,3 @@\n-old line\n+new line"
msg = BuildRequest(patch=patch_content, internal_patch_id="123")
# Mock the temporary file
temp_file_instance = MagicMock()
temp_file_instance.name = "/tmp/test_patch_file"
temp_file_mock.return_value.__enter__.return_value = temp_file_instance
result = self.builder_bot._apply_patch(task_mock, msg)
self.assertTrue(result)
temp_file_instance.write.assert_called_once_with(patch_content)
temp_file_instance.flush.assert_called_once()
task_mock.apply_patch_diff.assert_called_once_with(Path("/tmp/test_patch_file"))
@patch("tempfile.NamedTemporaryFile")
def test_apply_patch_with_patch_failure(self, temp_file_mock):
"""Test _apply_patch when patch is provided but fails to apply."""
task_mock = Mock()
task_mock.apply_patch_diff.return_value = False
patch_content = "--- a/file.c\n+++ b/file.c\n@@ -1,3 +1,3 @@\n-old line\n+new line"
msg = BuildRequest(patch=patch_content, internal_patch_id="123")
# Mock the temporary file
temp_file_instance = MagicMock()
temp_file_instance.name = "/tmp/test_patch_file"
temp_file_mock.return_value.__enter__.return_value = temp_file_instance
result = self.builder_bot._apply_patch(task_mock, msg)
self.assertFalse(result)
temp_file_instance.write.assert_called_once_with(patch_content)
temp_file_instance.flush.assert_called_once()
task_mock.apply_patch_diff.assert_called_once_with(Path("/tmp/test_patch_file"))
def test_serve_item_no_item_in_queue(self):
"""Test serve_item when there's no item in the queue."""
self.build_requests_queue_mock.pop.return_value = None
result = self.builder_bot.serve_item()
self.assertFalse(result)
self.build_requests_queue_mock.pop.assert_called_once()
@patch("buttercup.fuzzing_infra.builder_bot.ChallengeTask")
@patch("buttercup.fuzzing_infra.builder_bot.node_local.dir_to_remote_archive")
@patch("buttercup.fuzzing_infra.builder_bot.trace")
def test_serve_item_successful_build_with_caching(self, trace_mock, dir_to_remote_mock, challenge_task_mock):
"""Test serve_item with successful build when caching is enabled."""
# Setup queue item mock
queue_item_mock = Mock()
queue_item_mock.item_id = "test_item_id"
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="libfuzzer",
sanitizer="address",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=False,
patch="",
internal_patch_id="",
)
self.build_requests_queue_mock.pop.return_value = queue_item_mock
# Setup challenge task mocks
origin_task_mock = Mock()
origin_task_mock.task_meta.metadata = {"key": "value"} # Mock metadata as dict
task_mock = Mock()
task_mock.task_dir = Path("/tmp/test_task_dir")
# Mock build result
build_result_mock = Mock()
build_result_mock.success = True
task_mock.build_fuzzers_with_cache.return_value = build_result_mock
# Setup context manager for get_rw_copy
context_manager_mock = MagicMock()
context_manager_mock.__enter__.return_value = task_mock
context_manager_mock.__exit__.return_value = None
origin_task_mock.get_rw_copy.return_value = context_manager_mock
challenge_task_mock.return_value = origin_task_mock
# Setup tracer mock
tracer_mock = Mock()
span_mock = Mock()
span_context_manager_mock = MagicMock()
span_context_manager_mock.__enter__.return_value = span_mock
span_context_manager_mock.__exit__.return_value = None
tracer_mock.start_as_current_span.return_value = span_context_manager_mock
trace_mock.get_tracer.return_value = tracer_mock
# Call serve_item
result = self.builder_bot.serve_item()
# Verify behavior
self.assertTrue(result)
# Verify ChallengeTask was created with caching enabled
challenge_task_mock.assert_called_once_with(
Path("/path/to/task"),
python_path=self.python,
local_task_dir=Path("/path/to/task"),
)
# Verify build was called
task_mock.build_fuzzers_with_cache.assert_called_once_with(
engine="libfuzzer",
sanitizer="address",
pull_latest_base_image=self.allow_pull,
)
# Verify task was committed and archived
task_mock.commit.assert_called_once()
dir_to_remote_mock.assert_called_once_with(task_mock.task_dir)
# Verify build output was pushed with correct message
build_output_call = self.build_outputs_queue_mock.push.call_args[0][0]
self.assertIsInstance(build_output_call, BuildOutput)
self.assertEqual(build_output_call.task_id, "test_task")
self.assertEqual(build_output_call.engine, "libfuzzer")
self.assertEqual(build_output_call.sanitizer, "address")
self.assertEqual(build_output_call.build_type, BuildType.FUZZER)
self.assertEqual(build_output_call.task_dir, "/tmp/test_task_dir")
self.assertEqual(build_output_call.internal_patch_id, "")
# Verify item was acked
self.build_requests_queue_mock.ack_item.assert_called_once_with("test_item_id")
@patch("buttercup.fuzzing_infra.builder_bot.ChallengeTask")
@patch("buttercup.fuzzing_infra.builder_bot.node_local.dir_to_remote_archive")
@patch("buttercup.fuzzing_infra.builder_bot.trace")
def test_serve_item_successful_build_without_caching(self, trace_mock, dir_to_remote_mock, challenge_task_mock):
"""Test serve_item with successful build when caching is disabled."""
# Disable caching for this test
self.builder_bot.allow_caching = False
# Setup queue item mock
queue_item_mock = Mock()
queue_item_mock.item_id = "test_item_id"
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="libfuzzer",
sanitizer="address",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=False,
)
self.build_requests_queue_mock.pop.return_value = queue_item_mock
# Setup challenge task mocks
origin_task_mock = Mock()
origin_task_mock.task_meta.metadata = {"key": "value"} # Mock metadata as dict
task_mock = Mock()
task_mock.task_dir = Path("/tmp/test_task_dir")
# Mock build result
build_result_mock = Mock()
build_result_mock.success = True
task_mock.build_fuzzers_with_cache.return_value = build_result_mock
# Setup context manager for get_rw_copy
context_manager_mock = MagicMock()
context_manager_mock.__enter__.return_value = task_mock
context_manager_mock.__exit__.return_value = None
origin_task_mock.get_rw_copy.return_value = context_manager_mock
challenge_task_mock.return_value = origin_task_mock
# Setup tracer mock
tracer_mock = Mock()
span_mock = Mock()
span_context_manager_mock = MagicMock()
span_context_manager_mock.__enter__.return_value = span_mock
span_context_manager_mock.__exit__.return_value = None
tracer_mock.start_as_current_span.return_value = span_context_manager_mock
trace_mock.get_tracer.return_value = tracer_mock
# Call serve_item
result = self.builder_bot.serve_item()
# Verify behavior
self.assertTrue(result)
# Verify ChallengeTask was created without caching
challenge_task_mock.assert_called_once_with(Path("/path/to/task"), python_path=self.python)
# Verify build was called
task_mock.build_fuzzers_with_cache.assert_called_once_with(
engine="libfuzzer",
sanitizer="address",
pull_latest_base_image=self.allow_pull,
)
# Verify task was committed and archived
task_mock.commit.assert_called_once()
dir_to_remote_mock.assert_called_once_with(task_mock.task_dir)
# Verify build output was pushed
self.build_outputs_queue_mock.push.assert_called_once()
# Verify item was acked
self.build_requests_queue_mock.ack_item.assert_called_once_with("test_item_id")
@patch("buttercup.fuzzing_infra.builder_bot.ChallengeTask")
@patch("buttercup.fuzzing_infra.builder_bot.trace")
def test_serve_item_build_failure(self, trace_mock, challenge_task_mock):
"""Test serve_item when build fails."""
# Setup queue item mock
queue_item_mock = Mock()
queue_item_mock.item_id = "test_item_id"
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="libfuzzer",
sanitizer="address",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=False,
patch="",
internal_patch_id="",
)
self.build_requests_queue_mock.pop.return_value = queue_item_mock
# Setup challenge task mocks
origin_task_mock = Mock()
origin_task_mock.task_meta.metadata = {"key": "value"} # Mock metadata as dict
task_mock = Mock()
# Mock build result failure
build_result_mock = Mock()
build_result_mock.success = False
task_mock.build_fuzzers_with_cache.return_value = build_result_mock
# Setup context manager for get_rw_copy
context_manager_mock = MagicMock()
context_manager_mock.__enter__.return_value = task_mock
context_manager_mock.__exit__.return_value = None
origin_task_mock.get_rw_copy.return_value = context_manager_mock
challenge_task_mock.return_value = origin_task_mock
# Setup tracer mock
tracer_mock = Mock()
span_mock = Mock()
span_context_manager_mock = MagicMock()
span_context_manager_mock.__enter__.return_value = span_mock
span_context_manager_mock.__exit__.return_value = None
tracer_mock.start_as_current_span.return_value = span_context_manager_mock
trace_mock.get_tracer.return_value = tracer_mock
# Call serve_item
result = self.builder_bot.serve_item()
# Verify behavior
self.assertTrue(result) # Should still return True to continue processing
# Verify build was attempted
task_mock.build_fuzzers_with_cache.assert_called_once()
# Verify task was NOT committed (build failed)
task_mock.commit.assert_not_called()
# Verify build output was NOT pushed
self.build_outputs_queue_mock.push.assert_not_called()
# Verify span was marked as error
span_mock.set_status.assert_called()
@patch("buttercup.fuzzing_infra.builder_bot.ChallengeTask")
def test_serve_item_patch_application_failure(self, challenge_task_mock):
"""Test serve_item when patch application fails."""
# Setup queue item mock with a patch that will fail
queue_item_mock = Mock()
queue_item_mock.item_id = "test_item_id"
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="libfuzzer",
sanitizer="address",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=False,
patch="some patch content",
internal_patch_id="",
)
self.build_requests_queue_mock.pop.return_value = queue_item_mock
# Setup challenge task mocks
origin_task_mock = Mock()
origin_task_mock.task_meta.metadata = {"key": "value"} # Mock metadata as dict
task_mock = Mock()
task_mock.apply_patch_diff.return_value = False # Patch application fails
# Setup context manager for get_rw_copy
context_manager_mock = MagicMock()
context_manager_mock.__enter__.return_value = task_mock
context_manager_mock.__exit__.return_value = None
origin_task_mock.get_rw_copy.return_value = context_manager_mock
challenge_task_mock.return_value = origin_task_mock
# Mock times_delivered to return a value that exceeds max_tries
self.build_requests_queue_mock.times_delivered.return_value = 4
# Mock the _apply_patch method to return False
with patch.object(self.builder_bot, "_apply_patch", return_value=False):
# Call serve_item
result = self.builder_bot.serve_item()
# Verify behavior
self.assertTrue(result) # Should still return True to continue processing
# Verify build was NOT attempted (patch failed)
task_mock.build_fuzzers_with_cache.assert_not_called()
# Verify item was acked to avoid retrying forever
self.build_requests_queue_mock.ack_item.assert_called_once_with("test_item_id")
@patch("buttercup.fuzzing_infra.builder_bot.serve_loop")
def test_run(self, serve_loop_mock):
"""Test the run method."""
self.builder_bot.run()
# Verify serve_loop was called with correct parameters
serve_loop_mock.assert_called_once_with(self.builder_bot.serve_item, self.seconds_sleep)
@patch("buttercup.fuzzing_infra.builder_bot.ChallengeTask")
@patch("buttercup.fuzzing_infra.builder_bot.node_local.dir_to_remote_archive")
@patch("buttercup.fuzzing_infra.builder_bot.trace")
def test_serve_item_with_diff_and_patch(self, trace_mock, dir_to_remote_mock, challenge_task_mock):
"""Test serve_item with both diff and patch applied."""
# Setup queue item mock
queue_item_mock = Mock()
queue_item_mock.item_id = "test_item_id"
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="libfuzzer",
sanitizer="address",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=True,
patch="some patch content",
internal_patch_id="patch123",
)
self.build_requests_queue_mock.pop.return_value = queue_item_mock
# Setup challenge task mocks
origin_task_mock = Mock()
origin_task_mock.task_meta.metadata = {"key": "value"} # Mock metadata as dict
task_mock = Mock()
task_mock.task_dir = Path("/tmp/test_task_dir")
task_mock.apply_patch_diff.return_value = True # Both diff and patch succeed
# Mock build result
build_result_mock = Mock()
build_result_mock.success = True
task_mock.build_fuzzers_with_cache.return_value = build_result_mock
# Setup context manager for get_rw_copy
context_manager_mock = MagicMock()
context_manager_mock.__enter__.return_value = task_mock
context_manager_mock.__exit__.return_value = None
origin_task_mock.get_rw_copy.return_value = context_manager_mock
challenge_task_mock.return_value = origin_task_mock
# Setup tracer mock
tracer_mock = Mock()
span_mock = Mock()
span_context_manager_mock = MagicMock()
span_context_manager_mock.__enter__.return_value = span_mock
span_context_manager_mock.__exit__.return_value = None
tracer_mock.start_as_current_span.return_value = span_context_manager_mock
trace_mock.get_tracer.return_value = tracer_mock
# Call serve_item
result = self.builder_bot.serve_item()
# Verify behavior
self.assertTrue(result)
# Verify both diff and patch were applied
# apply_patch_diff should be called twice: once for diff, once for patch
self.assertEqual(task_mock.apply_patch_diff.call_count, 2)
# Verify build was successful
task_mock.build_fuzzers_with_cache.assert_called_once()
task_mock.commit.assert_called_once()
# Verify build output includes internal_patch_id
build_output_call = self.build_outputs_queue_mock.push.call_args[0][0]
self.assertEqual(build_output_call.internal_patch_id, "patch123")
@patch("buttercup.fuzzing_infra.builder_bot.ChallengeTask")
def test_max_tries_for_diff_and_patch(self, challenge_task_mock):
"""Test that max_tries is respected for both diff and patch application failures."""
# Test 1: Diff application failure - times_delivered < max_tries (should NOT ack)
queue_item_mock = Mock()
queue_item_mock.item_id = "test_item_id"
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="test_engine",
sanitizer="test_sanitizer",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=True,
patch="", # No patch for this test
internal_patch_id="",
)
self.build_requests_queue_mock.pop.return_value = queue_item_mock
# Setup challenge task mocks
origin_task_mock = Mock()
origin_task_mock.task_meta.metadata = {"key": "value"} # Mock metadata as dict
task_mock = Mock()
task_mock.task_dir = Path("/tmp/test_task_dir")
# Diff always fails
task_mock.apply_patch_diff.return_value = False
# Setup context manager for get_rw_copy
context_manager_mock = MagicMock()
context_manager_mock.__enter__.return_value = task_mock
context_manager_mock.__exit__.return_value = None
origin_task_mock.get_rw_copy.return_value = context_manager_mock
challenge_task_mock.return_value = origin_task_mock
# Mock times_delivered to simulate NOT exceeding max_tries (max_tries=3, so 2 should not ack)
self.build_requests_queue_mock.times_delivered.return_value = 2
# Call serve_item - should NOT ack because max tries not exceeded
result = self.builder_bot.serve_item()
self.assertTrue(result)
task_mock.apply_patch_diff.assert_called_once()
self.build_requests_queue_mock.ack_item.assert_not_called()
# Test 2: Diff application failure - times_delivered > max_tries (should ack)
self.build_requests_queue_mock.reset_mock()
task_mock.reset_mock()
# Mock times_delivered to simulate exceeding max_tries
self.build_requests_queue_mock.times_delivered.return_value = 4
# Call serve_item - should ack because max tries exceeded
result = self.builder_bot.serve_item()
self.assertTrue(result)
task_mock.apply_patch_diff.assert_called_once()
self.build_requests_queue_mock.ack_item.assert_called_once_with("test_item_id")
# Test 3: Patch application failure - times_delivered < max_tries (should NOT ack)
self.build_requests_queue_mock.reset_mock()
task_mock.reset_mock()
# New queue item with patch
queue_item_mock.deserialized = BuildRequest(
task_id="test_task",
engine="test_engine",
sanitizer="test_sanitizer",
build_type=BuildType.FUZZER,
task_dir="/path/to/task",
apply_diff=False, # No diff for this test
patch="some patch content",
internal_patch_id="patch123",
)
# Patch always fails
task_mock.apply_patch_diff.return_value = False
# Mock times_delivered to simulate NOT exceeding max_tries
self.build_requests_queue_mock.times_delivered.return_value = 2
# Call serve_item - should NOT ack because max tries not exceeded
result = self.builder_bot.serve_item()
self.assertTrue(result)
task_mock.apply_patch_diff.assert_called_once()
self.build_requests_queue_mock.ack_item.assert_not_called()
# Test 4: Patch application failure - times_delivered > max_tries (should ack)
self.build_requests_queue_mock.reset_mock()
task_mock.reset_mock()
# Mock times_delivered to simulate exceeding max_tries
self.build_requests_queue_mock.times_delivered.return_value = 4
# Call serve_item - should ack because max tries exceeded
result = self.builder_bot.serve_item()
self.assertTrue(result)
task_mock.apply_patch_diff.assert_called_once()
self.build_requests_queue_mock.ack_item.assert_called_once_with("test_item_id")
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink