From edcdb76f6b69b81eb8ea6acca4a993ed9957214b Mon Sep 17 00:00:00 2001 From: Joey Dreijer Date: Thu, 22 Oct 2026 19:58:25 +0200 Subject: [PATCH 1/3] fix: Make traversable edge property optional (for existing collectors) but default to None to discard value during conversion --- src/openhound/core/models/entries.py | 2 +- .../core/models/entries_dataclass.py | 6 +- tests/test_bhe_job_scheduling.py | 578 +----------------- tests/test_edge_properties.py | 10 + 4 files changed, 24 insertions(+), 572 deletions(-) create mode 100644 tests/test_edge_properties.py diff --git a/src/openhound/core/models/entries.py b/src/openhound/core/models/entries.py index f7e9658e..50485532 100644 --- a/src/openhound/core/models/entries.py +++ b/src/openhound/core/models/entries.py @@ -78,7 +78,7 @@ class EdgeProperties(BaseModel): model_config = ConfigDict(extra="allow") composed: bool = False - traversable: bool = False + traversable: bool | None = None class Edge(BaseModel): diff --git a/src/openhound/core/models/entries_dataclass.py b/src/openhound/core/models/entries_dataclass.py index 78d8c2be..6f5f1b90 100644 --- a/src/openhound/core/models/entries_dataclass.py +++ b/src/openhound/core/models/entries_dataclass.py @@ -76,7 +76,11 @@ class EdgeProperties: """Core model specifying required fields for an OpenGraph edge properties.""" composed: bool = field(default=False, kw_only=True) - traversable: bool = field(default=False, kw_only=True) + traversable: bool | None = field(default=None, kw_only=True) + + def __post_init__(self): + """Accept but discard traversable properties. None will be discarded during conversion.""" + self.traversable = None @dataclass diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 7e3746f3..dc88a9fd 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1,26 +1,14 @@ -import base64 import gzip -import hashlib import json -import logging from concurrent.futures import Future -from concurrent.futures.process import BrokenProcessPool from pathlib import Path -from types import SimpleNamespace from urllib.parse import urlsplit import pytest from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient -from openhound.core.clients import bloodhound, bloodhound_enterprise -from openhound.core.clients.bloodhound import BloodHoundHTTPError from openhound.core.clients.bloodhound_enterprise import JobStatus -from openhound.core.clients.models.jobs import ( - ManagementOperation, - ManagementOperationStatus, - ManagementOperationType, -) from openhound.core.models.graph import Graph from openhound.scheduler import service as scheduler_service from openhound.scheduler.service import ( @@ -31,7 +19,6 @@ ) TEST_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "jobs" -MANAGEMENT_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "management" def load_json(filename: str) -> dict: @@ -52,19 +39,9 @@ def mock_bloodhound_api(): app.state.job_ended = False app.state.end_payload = None app.state.start_payload = None - app.state.client_update_payload = None app.state.ingested_edges = 0 - app.state.management_operations = [] - app.state.operation_started = False - app.state.operation_ended = False - app.state.operation_start_payload = None - app.state.operation_end_payload = None - app.state.operation_completed_by_artifact_upload = False - app.state.bundle_content = None - app.state.artifact_create_payload = None - app.state.uploaded_parts = [] - app.state.artifact_completed = False app.state.ingested_nodes = 0 + app.state.ingested_edge_properties = [] @app.get("/api/v2/jobs/available") async def jobs_available(): @@ -92,105 +69,19 @@ async def end_job(body: dict): async def ingest(request: Request): body = await request.body() decompressed = gzip.decompress(body) + app.state.ingested_edge_properties.extend( + edge["properties"] for edge in json.loads(decompressed)["graph"]["edges"] + ) validate_graph = Graph.model_validate_json(decompressed) app.state.ingested_nodes += len(validate_graph.graph.nodes) app.state.ingested_edges += len(validate_graph.graph.edges) return {"status": "success"} - @app.put("/api/v2/clients/update") - async def update_client(body: dict): - app.state.client_update_payload = body - return {"status": "success"} - - @app.get("/api/v2/clients/management/available") - async def management_available(): - return {"data": app.state.management_operations} - - @app.post("/api/v2/clients/management/start") - async def start_operation(body: dict): - app.state.operation_started = True - app.state.operation_start_payload = body - return { - "data": { - "id": body["operation_id"], - "client_id": "client-123", - "artifact_id": None, - "type": "support_bundle", - "status": "running", - "created_at": "2026-01-01T00:00:00Z", - "updated_at": "2026-01-01T00:00:00Z", - "execution_time": "2026-01-01T00:00:00Z", - } - } - - @app.post("/api/v2/clients/management/artifacts") - async def create_artifact_upload(body: dict): - app.state.artifact_create_payload = body - return { - "data": { - "artifact_id": "artifact-123", - "client_id": "client-123", - "storage_key": "client-123--openhound-faker_support_bundle_2026-01-01_00-00-00.zip", - "status": "pending", - "part_size": body["part_size"], - "part_count": body["part_count"], - "missing_parts": list(range(1, body["part_count"] + 1)), - "management_operation": { - "id": body["operation_id"], - "client_id": "client-123", - "artifact_id": "artifact-123", - "type": "support_bundle", - "status": "running", - "requested_by_user_id": None, - "created_at": "2026-01-01T00:00:00Z", - "updated_at": "2026-01-01T00:00:00Z", - "started_at": "2026-01-01T00:00:00Z", - "completed_at": None, - "execution_time": "2026-01-01T00:00:00Z", - }, - } - } - - @app.post("/api/v2/clients/management/artifacts/{artifact_id}/parts/{part_number}") - async def upload_artifact_part( - artifact_id: str, part_number: int, request: Request - ): - content = await request.body() - checksum = base64.b64encode(hashlib.sha256(content).digest()).decode("ascii") - assert request.headers["content-digest"] == f"sha-256=:{checksum}:" - app.state.uploaded_parts.append((artifact_id, part_number, content)) - return Response(status_code=200) - - @app.post("/api/v2/clients/management/artifacts/{artifact_id}/complete") - async def complete_artifact_upload(artifact_id: str, body: dict): - app.state.artifact_completed = body["operation_id"] is not None - # BHE completes the associated management operation as part of this endpoint. - app.state.operation_completed_by_artifact_upload = True - return Response(status_code=204) - - @app.post("/api/v2/clients/management/end") - async def end_operation(body: dict): - app.state.operation_ended = True - app.state.operation_end_payload = body - return { - "data": { - "id": body["operation_id"], - "client_id": "client-123", - "artifact_id": "artifact-123", - "type": "support_bundle", - "status": body["status"], - "created_at": "2026-01-01T00:00:00Z", - "updated_at": "2026-01-01T00:00:00Z", - "execution_time": "2026-01-01T00:00:00Z", - } - } - return TestClient(app) @pytest.fixture def mock_service(mock_bloodhound_api, monkeypatch): - monkeypatch.setattr(bloodhound.openhound, "__version__", "0.3.0rc1") """Patches requests.requests so that our mocked BloodHound API will be used for testing the service. Args: @@ -216,8 +107,6 @@ def mock_request(method, url, **kwargs): return mock_bloodhound_api.get(path) if method.upper() == "POST": return mock_bloodhound_api.post(path, **kwargs) - if method.upper() == "PUT": - return mock_bloodhound_api.put(path, **kwargs) raise AssertionError(f"Unhandled method: {method}") @@ -228,102 +117,11 @@ def mock_request(method, url, **kwargs): bhe_uri="http://localhost:8000", token_key="test-key", token_id="test-id", + interval=1, collector_name="openhound-faker", ) -def test_client_update_sends_metadata(mock_service, mock_bloodhound_api, monkeypatch): - monkeypatch.setattr( - bloodhound_enterprise.socket, "gethostname", lambda: "test-host" - ) - monkeypatch.setattr( - bloodhound_enterprise.socket, - "gethostbyname", - lambda hostname: "192.0.2.10", - ) - - mock_service.client.update_client_metadata() - - assert mock_bloodhound_api.app.state.client_update_payload == { - "Address": "192.0.2.10", - "Hostname": "test-host", - "Version": "v0.3.0-rc1", - } - - -def test_request_forwards_connect_and_read_timeouts(monkeypatch): - captured = {} - - class Response: - status_code = 204 - text = "" - - def mock_request(**kwargs): - captured.update(kwargs) - return Response() - - monkeypatch.setattr(bloodhound.requests, "request", mock_request) - client = bloodhound.BloodHound(token_key="test-key", token_id="test-id") - - client.request(method="POST", path="/test", timeout=(1.5, 3.5)) - - assert captured["timeout"] == (1.5, 3.5) - - -def test_upload_artifact_part_uses_bounded_timeouts(mock_service, monkeypatch): - captured = {} - - def capture_request(*args, **kwargs): - captured.update(kwargs) - - monkeypatch.setattr(mock_service.client, "request", capture_request) - - mock_service.client.upload_artifact_part("artifact-123", 1, b"part") - - assert captured["timeout"] == ( - bloodhound_enterprise.SUPPORT_BUNDLE_CONNECT_TIMEOUT_SECONDS, - bloodhound_enterprise.SUPPORT_BUNDLE_READ_TIMEOUT_SECONDS, - ) - - -def test_client_update_uses_unknown_when_hostname_lookup_fails( - mock_service, mock_bloodhound_api, monkeypatch -): - def raise_error(): - raise OSError("hostname unavailable") - - monkeypatch.setattr(bloodhound_enterprise.socket, "gethostname", raise_error) - - mock_service.client.update_client_metadata() - - assert mock_bloodhound_api.app.state.client_update_payload == { - "Address": "unknown", - "Hostname": "unknown", - "Version": "v0.3.0-rc1", - } - - -def test_client_update_uses_unknown_when_ip_lookup_fails( - mock_service, mock_bloodhound_api, monkeypatch -): - monkeypatch.setattr( - bloodhound_enterprise.socket, "gethostname", lambda: "test-host" - ) - - def raise_error(hostname: str): - raise OSError(f"{hostname} unavailable") - - monkeypatch.setattr(bloodhound_enterprise.socket, "gethostbyname", raise_error) - - mock_service.client.update_client_metadata() - - assert mock_bloodhound_api.app.state.client_update_payload == { - "Address": "unknown", - "Hostname": "test-host", - "Version": "v0.3.0-rc1", - } - - def test_jobs_starts_new_job(mock_service, mock_bloodhound_api): """Runs _check_jobs and checks if the new job is started when available.""" @@ -407,110 +205,6 @@ def test_poll_missing_extension(mock_service, mock_bloodhound_api): } -def test_poll_recovers_from_broken_process_pool(mock_service, mock_bloodhound_api): - """A BrokenProcessPool surfaced via future.result() should fail the job, clear state, and rebuild the executor.""" - mock_bloodhound_api.app.state.job_started = True - future = Future() - future.set_exception(BrokenProcessPool("worker died")) - mock_service.future = future - mock_service.job_running = 123 - original_executor = mock_service.executor - - mock_service._poll() - - assert mock_service.future is None - assert mock_service.job_running is None - assert mock_service.executor is not original_executor - assert mock_bloodhound_api.app.state.job_ended is True - assert mock_bloodhound_api.app.state.end_payload == { - "status": JobStatus.FAILED.value, - "message": "Collection worker for 'openhound-faker' was terminated abruptly", - } - - -def test_start_job_recovers_when_submit_raises_broken_pool( - mock_service, mock_bloodhound_api, monkeypatch -): - """If executor.submit raises BrokenProcessPool after the BHE job was started, the job should be ended FAILED, state cleared, and the executor rebuilt.""" - - def broken_submit(*args, **kwargs): - raise BrokenProcessPool("worker died before submit") - - monkeypatch.setattr(mock_service.executor, "submit", broken_submit) - original_executor = mock_service.executor - - mock_service._poll() - - assert mock_service.future is None - assert mock_service.job_running is None - assert mock_service.executor is not original_executor - assert mock_bloodhound_api.app.state.job_started is True - assert mock_bloodhound_api.app.state.job_ended is True - assert mock_bloodhound_api.app.state.end_payload == { - "status": JobStatus.FAILED.value, - "message": "Failed to start collector 'openhound-faker': worker pool was broken", - } - - -def test_checkin_calls_jobs_current_when_job_running(mock_service, monkeypatch): - """_poll() should call jobs_current via the else-branch check-in when a job is running.""" - # Simulate a job in progress with no completed future — skips the completion handler, - # reaches the else-branch, and triggers jobs_current as a check-in heartbeat. - mock_service.job_running = 123 - mock_service.future = None # no completed future to handle - called = [] - - def fake_jobs_current(self): - called.append(True) - - monkeypatch.setattr( - mock_service.client.__class__, "jobs_current", property(fake_jobs_current) - ) - - mock_service._poll() - - assert len(called) == 1 - - -def test_checkin_noop_when_no_job_running(mock_service, monkeypatch): - """_poll() should not call jobs_current via the else-branch check-in when no job is running.""" - # When idle (job_running is None), _poll() takes the if-branch and calls check_jobs() - # instead of the else-branch check-in. jobs_current should never be touched. - assert mock_service.job_running is None - mock_service.future = None - called = [] - - def fake_jobs_current(self): - called.append(True) - - monkeypatch.setattr( - mock_service.client.__class__, "jobs_current", property(fake_jobs_current) - ) - # Stub check_jobs so _poll doesn't try to start a job; we only care the else-branch doesn't fire - monkeypatch.setattr(mock_service, "check_jobs", lambda: None) - - mock_service._poll() - - assert len(called) == 0 - - -def test_checkin_swallows_exception(mock_service, monkeypatch): - """_poll() should swallow exceptions raised by jobs_current in the check-in else-branch.""" - # A transient BHE error during check-in must not crash the service loop. - mock_service.job_running = 123 - mock_service.future = None # no completed future to handle - - def raise_error(self): - raise RuntimeError("BHE unreachable") - - monkeypatch.setattr( - mock_service.client.__class__, "jobs_current", property(raise_error) - ) - - # Should not raise — _poll's except block absorbs the error - mock_service._poll() - - def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypatch): """Run the DLT pipeline with the openhound-faker collector + check the amount of ingested nodes + edges""" monkeypatch.setenv( @@ -524,263 +218,7 @@ def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypat assert result.job_id == 123 assert mock_bloodhound_api.app.state.ingested_nodes == 1000 assert mock_bloodhound_api.app.state.ingested_edges == 10000 - - -def _support_bundle_operation() -> dict: - return json.loads( - (MANAGEMENT_DATA_DIR / "management_available_with_operation.json").read_text() - )["data"][0] - - -def test_check_management_returns_support_bundle_operation( - mock_service, mock_bloodhound_api -): - mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] - - operation = mock_service.check_management() - - assert operation is not None - assert operation.type is ManagementOperationType.SUPPORT_BUNDLE - - -def test_check_management_ignores_non_queued_operations( - mock_service, mock_bloodhound_api -): - operation = _support_bundle_operation() - operation["status"] = ManagementOperationStatus.RUNNING.value - mock_bloodhound_api.app.state.management_operations = [operation] - - assert mock_service.check_management() is None - - -def test_poll_prioritizes_management_over_a_new_job( - mock_service, mock_bloodhound_api, monkeypatch -): - mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] - sent = [] - monkeypatch.setattr(mock_service, "_send_support_bundle", sent.append) - - mock_service._poll() - - assert len(sent) == 1 - assert mock_bloodhound_api.app.state.job_started is False - - -def test_poll_starts_job_when_no_management_work( - mock_service, mock_bloodhound_api, monkeypatch -): - submitted = Future() - monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) - - mock_service._poll() - - assert mock_bloodhound_api.app.state.job_started is True - - -def test_poll_still_checks_jobs_when_management_endpoint_fails( - mock_service, mock_bloodhound_api, monkeypatch -): - submitted = Future() - monkeypatch.setattr(mock_service, "check_management", lambda: 1 / 0) - monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) - - mock_service._poll() - - assert mock_bloodhound_api.app.state.job_started is True - - -def test_poll_does_not_start_a_job_when_management_work_fails( - mock_service, mock_bloodhound_api, monkeypatch -): - mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] - - def fail(operation): - raise RuntimeError("upload failed") - - monkeypatch.setattr(mock_service, "_send_support_bundle", fail) - - mock_service._poll() - - assert mock_bloodhound_api.app.state.job_started is False - - -def test_send_support_bundle_claims_uploads_completes_and_cleans_up( - mock_service, mock_bloodhound_api, tmp_path, monkeypatch -): - log = tmp_path / "openhound.log" - log.write_text("support log") - mock_service.log_base_path = tmp_path - created = [] - - from openhound.scheduler import service as scheduler_service - - original_create = scheduler_service.create_support_bundle - - def capture_bundle(*args): - bundle = original_create(*args) - created.append(bundle) - return bundle - - monkeypatch.setattr(scheduler_service, "create_support_bundle", capture_bundle) - operation = ManagementOperation.model_validate(_support_bundle_operation()) - - mock_service._send_support_bundle(operation) - - assert mock_bloodhound_api.app.state.operation_start_payload == { - "operation_id": operation.id - } - assert ( - mock_bloodhound_api.app.state.artifact_create_payload["operation_id"] - == operation.id - ) - assert mock_bloodhound_api.app.state.uploaded_parts - assert mock_bloodhound_api.app.state.artifact_completed is True - assert mock_bloodhound_api.app.state.operation_completed_by_artifact_upload is True - assert mock_bloodhound_api.app.state.operation_end_payload == { - "operation_id": operation.id, - "status": ManagementOperationStatus.SUCCEEDED.value, - } - assert created and not created[0].exists() - assert not created[0].parent.exists() - - -def test_create_artifact_upload_preserves_entire_create_response( - mock_service, tmp_path -): - bundle = tmp_path / "support-bundle.zip" - bundle.write_bytes(b"support bundle") - - session = mock_service.client.create_artifact_upload("operation-123", bundle) - - assert session.artifact_id == "artifact-123" - assert session.client_id == "client-123" - assert session.storage_key.endswith("support_bundle_2026-01-01_00-00-00.zip") - assert session.status == "pending" - assert session.missing_parts == [1] - assert session.management_operation.id == "operation-123" - assert session.management_operation.artifact_id == session.artifact_id - - -@pytest.mark.parametrize( - "failure_point", - [ - "start_operation", - "create_support_bundle", - "create_artifact_upload", - "upload_artifact_part", - "complete_artifact_upload", - ], -) -def test_send_support_bundle_marks_operation_failed_for_each_lifecycle_failure( - mock_service, mock_bloodhound_api, monkeypatch, caplog, failure_point -): - operation = ManagementOperation.model_validate(_support_bundle_operation()) - - def fail(*args, **kwargs): - raise RuntimeError(f"{failure_point} failed") - - if failure_point == "create_support_bundle": - monkeypatch.setattr(scheduler_service, "create_support_bundle", fail) - else: - monkeypatch.setattr(mock_service.client, failure_point, fail) - - with pytest.raises(RuntimeError, match=f"{failure_point} failed"): - mock_service._send_support_bundle(operation) - - assert mock_bloodhound_api.app.state.operation_end_payload == { - "operation_id": operation.id, - "status": ManagementOperationStatus.FAILED.value, - } - assert "Support bundle operation" in caplog.text - - -def test_send_support_bundle_retries_transient_part_upload_failure( - mock_service, monkeypatch, tmp_path -): - log = tmp_path / "openhound.log" - log.write_text("support log") - mock_service.log_base_path = tmp_path - operation = ManagementOperation.model_validate(_support_bundle_operation()) - original_request = mock_service.client.request - attempts = 0 - delays = [] - - def flaky_request(method, path, **kwargs): - nonlocal attempts - if "/parts/" in path: - attempts += 1 - if attempts < 3: - raise BloodHoundHTTPError("temporary failure", 503) - return original_request(method, path, **kwargs) - - monkeypatch.setattr(mock_service.client, "request", flaky_request) - monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) - - mock_service._send_support_bundle(operation) - - assert attempts == 3 - assert delays == [2, 2] - - -def test_send_support_bundle_fails_after_transient_retries_are_exhausted( - mock_service, mock_bloodhound_api, monkeypatch, tmp_path -): - log = tmp_path / "openhound.log" - log.write_text("support log") - mock_service.log_base_path = tmp_path - operation = ManagementOperation.model_validate(_support_bundle_operation()) - original_request = mock_service.client.request - attempts = 0 - delays = [] - - def unavailable_part_upload(method, path, **kwargs): - nonlocal attempts - if "/parts/" in path: - attempts += 1 - raise BloodHoundHTTPError("temporarily unavailable", 503) - return original_request(method, path, **kwargs) - - monkeypatch.setattr(mock_service.client, "request", unavailable_part_upload) - monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) - - with pytest.raises(BloodHoundHTTPError): - mock_service._send_support_bundle(operation) - - assert attempts == 4 - assert delays == [2, 2, 2] - assert mock_bloodhound_api.app.state.operation_end_payload == { - "operation_id": operation.id, - "status": ManagementOperationStatus.FAILED.value, - } - - -def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, caplog): - collector = SimpleNamespace( - name="example", - package_version="2.3.4", - metadata=SimpleNamespace(version="1.2.3"), - ) - monkeypatch.setattr( - scheduler_service.CollectorManager, - "from_entrypoint", - lambda: SimpleNamespace(collectors=[collector]), - ) - monkeypatch.setattr( - scheduler_service.dataflow, "pipeline", lambda extension: {"collect": []} + assert all( + "traversable" not in properties + for properties in mock_bloodhound_api.app.state.ingested_edge_properties ) - monkeypatch.setattr(scheduler_service.openhound, "__version__", "4.5.6") - - with caplog.at_level(logging.INFO, logger="openhound.scheduler.service"): - _subprocess_collect("example", 42) - - collection_records = [ - record - for record in caplog.records - if record.getMessage().startswith(("Subprocess running collection", "Collection for job")) - ] - assert len(collection_records) == 2 - for record in collection_records: - assert record.collector_extension == "example" - assert record.collector_extension_version == "2.3.4" - assert record.openhound_version == "4.5.6" - assert record.job_id == 42 diff --git a/tests/test_edge_properties.py b/tests/test_edge_properties.py new file mode 100644 index 00000000..2b30a7dd --- /dev/null +++ b/tests/test_edge_properties.py @@ -0,0 +1,10 @@ +import pytest + +from openhound.core.models.entries_dataclass import EdgeProperties + + +@pytest.mark.parametrize("traversable", [True, False]) +def test_edge_properties_accept_and_discard_legacy_traversable(traversable): + properties = EdgeProperties(traversable=traversable) + + assert properties.traversable is None From 4866cec5ce2b5bc3b90ef59891e414b344be95be Mon Sep 17 00:00:00 2001 From: Joey Dreijer Date: Thu, 22 Oct 2026 20:03:08 +0200 Subject: [PATCH 2/3] fix: Merged test with state of main --- tests/test_bhe_job_scheduling.py | 573 ++++++++++++++++++++++++++++++- 1 file changed, 572 insertions(+), 1 deletion(-) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index dc88a9fd..69a19e12 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1,14 +1,26 @@ +import base64 import gzip +import hashlib import json +import logging from concurrent.futures import Future +from concurrent.futures.process import BrokenProcessPool from pathlib import Path +from types import SimpleNamespace from urllib.parse import urlsplit import pytest from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient +from openhound.core.clients import bloodhound, bloodhound_enterprise +from openhound.core.clients.bloodhound import BloodHoundHTTPError from openhound.core.clients.bloodhound_enterprise import JobStatus +from openhound.core.clients.models.jobs import ( + ManagementOperation, + ManagementOperationStatus, + ManagementOperationType, +) from openhound.core.models.graph import Graph from openhound.scheduler import service as scheduler_service from openhound.scheduler.service import ( @@ -19,6 +31,7 @@ ) TEST_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "jobs" +MANAGEMENT_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "management" def load_json(filename: str) -> dict: @@ -39,7 +52,18 @@ def mock_bloodhound_api(): app.state.job_ended = False app.state.end_payload = None app.state.start_payload = None + app.state.client_update_payload = None app.state.ingested_edges = 0 + app.state.management_operations = [] + app.state.operation_started = False + app.state.operation_ended = False + app.state.operation_start_payload = None + app.state.operation_end_payload = None + app.state.operation_completed_by_artifact_upload = False + app.state.bundle_content = None + app.state.artifact_create_payload = None + app.state.uploaded_parts = [] + app.state.artifact_completed = False app.state.ingested_nodes = 0 app.state.ingested_edge_properties = [] @@ -77,11 +101,100 @@ async def ingest(request: Request): app.state.ingested_edges += len(validate_graph.graph.edges) return {"status": "success"} + @app.put("/api/v2/clients/update") + async def update_client(body: dict): + app.state.client_update_payload = body + return {"status": "success"} + + @app.get("/api/v2/clients/management/available") + async def management_available(): + return {"data": app.state.management_operations} + + @app.post("/api/v2/clients/management/start") + async def start_operation(body: dict): + app.state.operation_started = True + app.state.operation_start_payload = body + return { + "data": { + "id": body["operation_id"], + "client_id": "client-123", + "artifact_id": None, + "type": "support_bundle", + "status": "running", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + "execution_time": "2026-01-01T00:00:00Z", + } + } + + @app.post("/api/v2/clients/management/artifacts") + async def create_artifact_upload(body: dict): + app.state.artifact_create_payload = body + return { + "data": { + "artifact_id": "artifact-123", + "client_id": "client-123", + "storage_key": "client-123--openhound-faker_support_bundle_2026-01-01_00-00-00.zip", + "status": "pending", + "part_size": body["part_size"], + "part_count": body["part_count"], + "missing_parts": list(range(1, body["part_count"] + 1)), + "management_operation": { + "id": body["operation_id"], + "client_id": "client-123", + "artifact_id": "artifact-123", + "type": "support_bundle", + "status": "running", + "requested_by_user_id": None, + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + "started_at": "2026-01-01T00:00:00Z", + "completed_at": None, + "execution_time": "2026-01-01T00:00:00Z", + }, + } + } + + @app.post("/api/v2/clients/management/artifacts/{artifact_id}/parts/{part_number}") + async def upload_artifact_part( + artifact_id: str, part_number: int, request: Request + ): + content = await request.body() + checksum = base64.b64encode(hashlib.sha256(content).digest()).decode("ascii") + assert request.headers["content-digest"] == f"sha-256=:{checksum}:" + app.state.uploaded_parts.append((artifact_id, part_number, content)) + return Response(status_code=200) + + @app.post("/api/v2/clients/management/artifacts/{artifact_id}/complete") + async def complete_artifact_upload(artifact_id: str, body: dict): + app.state.artifact_completed = body["operation_id"] is not None + # BHE completes the associated management operation as part of this endpoint. + app.state.operation_completed_by_artifact_upload = True + return Response(status_code=204) + + @app.post("/api/v2/clients/management/end") + async def end_operation(body: dict): + app.state.operation_ended = True + app.state.operation_end_payload = body + return { + "data": { + "id": body["operation_id"], + "client_id": "client-123", + "artifact_id": "artifact-123", + "type": "support_bundle", + "status": body["status"], + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + "execution_time": "2026-01-01T00:00:00Z", + } + } + return TestClient(app) @pytest.fixture def mock_service(mock_bloodhound_api, monkeypatch): + monkeypatch.setattr(bloodhound.openhound, "__version__", "0.3.0rc1") """Patches requests.requests so that our mocked BloodHound API will be used for testing the service. Args: @@ -107,6 +220,8 @@ def mock_request(method, url, **kwargs): return mock_bloodhound_api.get(path) if method.upper() == "POST": return mock_bloodhound_api.post(path, **kwargs) + if method.upper() == "PUT": + return mock_bloodhound_api.put(path, **kwargs) raise AssertionError(f"Unhandled method: {method}") @@ -117,11 +232,102 @@ def mock_request(method, url, **kwargs): bhe_uri="http://localhost:8000", token_key="test-key", token_id="test-id", - interval=1, collector_name="openhound-faker", ) +def test_client_update_sends_metadata(mock_service, mock_bloodhound_api, monkeypatch): + monkeypatch.setattr( + bloodhound_enterprise.socket, "gethostname", lambda: "test-host" + ) + monkeypatch.setattr( + bloodhound_enterprise.socket, + "gethostbyname", + lambda hostname: "192.0.2.10", + ) + + mock_service.client.update_client_metadata() + + assert mock_bloodhound_api.app.state.client_update_payload == { + "Address": "192.0.2.10", + "Hostname": "test-host", + "Version": "v0.3.0-rc1", + } + + +def test_request_forwards_connect_and_read_timeouts(monkeypatch): + captured = {} + + class Response: + status_code = 204 + text = "" + + def mock_request(**kwargs): + captured.update(kwargs) + return Response() + + monkeypatch.setattr(bloodhound.requests, "request", mock_request) + client = bloodhound.BloodHound(token_key="test-key", token_id="test-id") + + client.request(method="POST", path="/test", timeout=(1.5, 3.5)) + + assert captured["timeout"] == (1.5, 3.5) + + +def test_upload_artifact_part_uses_bounded_timeouts(mock_service, monkeypatch): + captured = {} + + def capture_request(*args, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(mock_service.client, "request", capture_request) + + mock_service.client.upload_artifact_part("artifact-123", 1, b"part") + + assert captured["timeout"] == ( + bloodhound_enterprise.SUPPORT_BUNDLE_CONNECT_TIMEOUT_SECONDS, + bloodhound_enterprise.SUPPORT_BUNDLE_READ_TIMEOUT_SECONDS, + ) + + +def test_client_update_uses_unknown_when_hostname_lookup_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + def raise_error(): + raise OSError("hostname unavailable") + + monkeypatch.setattr(bloodhound_enterprise.socket, "gethostname", raise_error) + + mock_service.client.update_client_metadata() + + assert mock_bloodhound_api.app.state.client_update_payload == { + "Address": "unknown", + "Hostname": "unknown", + "Version": "v0.3.0-rc1", + } + + +def test_client_update_uses_unknown_when_ip_lookup_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + monkeypatch.setattr( + bloodhound_enterprise.socket, "gethostname", lambda: "test-host" + ) + + def raise_error(hostname: str): + raise OSError(f"{hostname} unavailable") + + monkeypatch.setattr(bloodhound_enterprise.socket, "gethostbyname", raise_error) + + mock_service.client.update_client_metadata() + + assert mock_bloodhound_api.app.state.client_update_payload == { + "Address": "unknown", + "Hostname": "test-host", + "Version": "v0.3.0-rc1", + } + + def test_jobs_starts_new_job(mock_service, mock_bloodhound_api): """Runs _check_jobs and checks if the new job is started when available.""" @@ -205,6 +411,110 @@ def test_poll_missing_extension(mock_service, mock_bloodhound_api): } +def test_poll_recovers_from_broken_process_pool(mock_service, mock_bloodhound_api): + """A BrokenProcessPool surfaced via future.result() should fail the job, clear state, and rebuild the executor.""" + mock_bloodhound_api.app.state.job_started = True + future = Future() + future.set_exception(BrokenProcessPool("worker died")) + mock_service.future = future + mock_service.job_running = 123 + original_executor = mock_service.executor + + mock_service._poll() + + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.executor is not original_executor + assert mock_bloodhound_api.app.state.job_ended is True + assert mock_bloodhound_api.app.state.end_payload == { + "status": JobStatus.FAILED.value, + "message": "Collection worker for 'openhound-faker' was terminated abruptly", + } + + +def test_start_job_recovers_when_submit_raises_broken_pool( + mock_service, mock_bloodhound_api, monkeypatch +): + """If executor.submit raises BrokenProcessPool after the BHE job was started, the job should be ended FAILED, state cleared, and the executor rebuilt.""" + + def broken_submit(*args, **kwargs): + raise BrokenProcessPool("worker died before submit") + + monkeypatch.setattr(mock_service.executor, "submit", broken_submit) + original_executor = mock_service.executor + + mock_service._poll() + + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.executor is not original_executor + assert mock_bloodhound_api.app.state.job_started is True + assert mock_bloodhound_api.app.state.job_ended is True + assert mock_bloodhound_api.app.state.end_payload == { + "status": JobStatus.FAILED.value, + "message": "Failed to start collector 'openhound-faker': worker pool was broken", + } + + +def test_checkin_calls_jobs_current_when_job_running(mock_service, monkeypatch): + """_poll() should call jobs_current via the else-branch check-in when a job is running.""" + # Simulate a job in progress with no completed future — skips the completion handler, + # reaches the else-branch, and triggers jobs_current as a check-in heartbeat. + mock_service.job_running = 123 + mock_service.future = None # no completed future to handle + called = [] + + def fake_jobs_current(self): + called.append(True) + + monkeypatch.setattr( + mock_service.client.__class__, "jobs_current", property(fake_jobs_current) + ) + + mock_service._poll() + + assert len(called) == 1 + + +def test_checkin_noop_when_no_job_running(mock_service, monkeypatch): + """_poll() should not call jobs_current via the else-branch check-in when no job is running.""" + # When idle (job_running is None), _poll() takes the if-branch and calls check_jobs() + # instead of the else-branch check-in. jobs_current should never be touched. + assert mock_service.job_running is None + mock_service.future = None + called = [] + + def fake_jobs_current(self): + called.append(True) + + monkeypatch.setattr( + mock_service.client.__class__, "jobs_current", property(fake_jobs_current) + ) + # Stub check_jobs so _poll doesn't try to start a job; we only care the else-branch doesn't fire + monkeypatch.setattr(mock_service, "check_jobs", lambda: None) + + mock_service._poll() + + assert len(called) == 0 + + +def test_checkin_swallows_exception(mock_service, monkeypatch): + """_poll() should swallow exceptions raised by jobs_current in the check-in else-branch.""" + # A transient BHE error during check-in must not crash the service loop. + mock_service.job_running = 123 + mock_service.future = None # no completed future to handle + + def raise_error(self): + raise RuntimeError("BHE unreachable") + + monkeypatch.setattr( + mock_service.client.__class__, "jobs_current", property(raise_error) + ) + + # Should not raise — _poll's except block absorbs the error + mock_service._poll() + + def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypatch): """Run the DLT pipeline with the openhound-faker collector + check the amount of ingested nodes + edges""" monkeypatch.setenv( @@ -218,6 +528,267 @@ def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypat assert result.job_id == 123 assert mock_bloodhound_api.app.state.ingested_nodes == 1000 assert mock_bloodhound_api.app.state.ingested_edges == 10000 + + +def _support_bundle_operation() -> dict: + return json.loads( + (MANAGEMENT_DATA_DIR / "management_available_with_operation.json").read_text() + )["data"][0] + + +def test_check_management_returns_support_bundle_operation( + mock_service, mock_bloodhound_api +): + mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] + + operation = mock_service.check_management() + + assert operation is not None + assert operation.type is ManagementOperationType.SUPPORT_BUNDLE + + +def test_check_management_ignores_non_queued_operations( + mock_service, mock_bloodhound_api +): + operation = _support_bundle_operation() + operation["status"] = ManagementOperationStatus.RUNNING.value + mock_bloodhound_api.app.state.management_operations = [operation] + + assert mock_service.check_management() is None + + +def test_poll_prioritizes_management_over_a_new_job( + mock_service, mock_bloodhound_api, monkeypatch +): + mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] + sent = [] + monkeypatch.setattr(mock_service, "_send_support_bundle", sent.append) + + mock_service._poll() + + assert len(sent) == 1 + assert mock_bloodhound_api.app.state.job_started is False + + +def test_poll_starts_job_when_no_management_work( + mock_service, mock_bloodhound_api, monkeypatch +): + submitted = Future() + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + + mock_service._poll() + + assert mock_bloodhound_api.app.state.job_started is True + + +def test_poll_still_checks_jobs_when_management_endpoint_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + submitted = Future() + monkeypatch.setattr(mock_service, "check_management", lambda: 1 / 0) + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + + mock_service._poll() + + assert mock_bloodhound_api.app.state.job_started is True + + +def test_poll_does_not_start_a_job_when_management_work_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] + + def fail(operation): + raise RuntimeError("upload failed") + + monkeypatch.setattr(mock_service, "_send_support_bundle", fail) + + mock_service._poll() + + assert mock_bloodhound_api.app.state.job_started is False + + +def test_send_support_bundle_claims_uploads_completes_and_cleans_up( + mock_service, mock_bloodhound_api, tmp_path, monkeypatch +): + log = tmp_path / "openhound.log" + log.write_text("support log") + mock_service.log_base_path = tmp_path + created = [] + + from openhound.scheduler import service as scheduler_service + + original_create = scheduler_service.create_support_bundle + + def capture_bundle(*args): + bundle = original_create(*args) + created.append(bundle) + return bundle + + monkeypatch.setattr(scheduler_service, "create_support_bundle", capture_bundle) + operation = ManagementOperation.model_validate(_support_bundle_operation()) + + mock_service._send_support_bundle(operation) + + assert mock_bloodhound_api.app.state.operation_start_payload == { + "operation_id": operation.id + } + assert ( + mock_bloodhound_api.app.state.artifact_create_payload["operation_id"] + == operation.id + ) + assert mock_bloodhound_api.app.state.uploaded_parts + assert mock_bloodhound_api.app.state.artifact_completed is True + assert mock_bloodhound_api.app.state.operation_completed_by_artifact_upload is True + assert mock_bloodhound_api.app.state.operation_end_payload == { + "operation_id": operation.id, + "status": ManagementOperationStatus.SUCCEEDED.value, + } + assert created and not created[0].exists() + assert not created[0].parent.exists() + + +def test_create_artifact_upload_preserves_entire_create_response( + mock_service, tmp_path +): + bundle = tmp_path / "support-bundle.zip" + bundle.write_bytes(b"support bundle") + + session = mock_service.client.create_artifact_upload("operation-123", bundle) + + assert session.artifact_id == "artifact-123" + assert session.client_id == "client-123" + assert session.storage_key.endswith("support_bundle_2026-01-01_00-00-00.zip") + assert session.status == "pending" + assert session.missing_parts == [1] + assert session.management_operation.id == "operation-123" + assert session.management_operation.artifact_id == session.artifact_id + + +@pytest.mark.parametrize( + "failure_point", + [ + "start_operation", + "create_support_bundle", + "create_artifact_upload", + "upload_artifact_part", + "complete_artifact_upload", + ], +) +def test_send_support_bundle_marks_operation_failed_for_each_lifecycle_failure( + mock_service, mock_bloodhound_api, monkeypatch, caplog, failure_point +): + operation = ManagementOperation.model_validate(_support_bundle_operation()) + + def fail(*args, **kwargs): + raise RuntimeError(f"{failure_point} failed") + + if failure_point == "create_support_bundle": + monkeypatch.setattr(scheduler_service, "create_support_bundle", fail) + else: + monkeypatch.setattr(mock_service.client, failure_point, fail) + + with pytest.raises(RuntimeError, match=f"{failure_point} failed"): + mock_service._send_support_bundle(operation) + + assert mock_bloodhound_api.app.state.operation_end_payload == { + "operation_id": operation.id, + "status": ManagementOperationStatus.FAILED.value, + } + assert "Support bundle operation" in caplog.text + + +def test_send_support_bundle_retries_transient_part_upload_failure( + mock_service, monkeypatch, tmp_path +): + log = tmp_path / "openhound.log" + log.write_text("support log") + mock_service.log_base_path = tmp_path + operation = ManagementOperation.model_validate(_support_bundle_operation()) + original_request = mock_service.client.request + attempts = 0 + delays = [] + + def flaky_request(method, path, **kwargs): + nonlocal attempts + if "/parts/" in path: + attempts += 1 + if attempts < 3: + raise BloodHoundHTTPError("temporary failure", 503) + return original_request(method, path, **kwargs) + + monkeypatch.setattr(mock_service.client, "request", flaky_request) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + mock_service._send_support_bundle(operation) + + assert attempts == 3 + assert delays == [2, 2] + + +def test_send_support_bundle_fails_after_transient_retries_are_exhausted( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path +): + log = tmp_path / "openhound.log" + log.write_text("support log") + mock_service.log_base_path = tmp_path + operation = ManagementOperation.model_validate(_support_bundle_operation()) + original_request = mock_service.client.request + attempts = 0 + delays = [] + + def unavailable_part_upload(method, path, **kwargs): + nonlocal attempts + if "/parts/" in path: + attempts += 1 + raise BloodHoundHTTPError("temporarily unavailable", 503) + return original_request(method, path, **kwargs) + + monkeypatch.setattr(mock_service.client, "request", unavailable_part_upload) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with pytest.raises(BloodHoundHTTPError): + mock_service._send_support_bundle(operation) + + assert attempts == 4 + assert delays == [2, 2, 2] + assert mock_bloodhound_api.app.state.operation_end_payload == { + "operation_id": operation.id, + "status": ManagementOperationStatus.FAILED.value, + } + + +def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, caplog): + collector = SimpleNamespace( + name="example", + package_version="2.3.4", + metadata=SimpleNamespace(version="1.2.3"), + ) + monkeypatch.setattr( + scheduler_service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + monkeypatch.setattr( + scheduler_service.dataflow, "pipeline", lambda extension: {"collect": []} + ) + monkeypatch.setattr(scheduler_service.openhound, "__version__", "4.5.6") + + with caplog.at_level(logging.INFO, logger="openhound.scheduler.service"): + _subprocess_collect("example", 42) + + collection_records = [ + record + for record in caplog.records + if record.getMessage().startswith(("Subprocess running collection", "Collection for job")) + ] + assert len(collection_records) == 2 + for record in collection_records: + assert record.collector_extension == "example" + assert record.collector_extension_version == "2.3.4" + assert record.openhound_version == "4.5.6" + assert record.job_id == 42 + assert all( "traversable" not in properties for properties in mock_bloodhound_api.app.state.ingested_edge_properties From eb215c6a2a9f7ba85b51e4ba794ec883661d8b42 Mon Sep 17 00:00:00 2001 From: Joey Dreijer Date: Thu, 22 Oct 2026 20:05:20 +0200 Subject: [PATCH 3/3] fix: moved traversable check to correct test --- tests/test_bhe_job_scheduling.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 69a19e12..b5287f84 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -528,6 +528,10 @@ def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypat assert result.job_id == 123 assert mock_bloodhound_api.app.state.ingested_nodes == 1000 assert mock_bloodhound_api.app.state.ingested_edges == 10000 + assert all( + "traversable" not in properties + for properties in mock_bloodhound_api.app.state.ingested_edge_properties + ) def _support_bundle_operation() -> dict: @@ -788,8 +792,3 @@ def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, c assert record.collector_extension_version == "2.3.4" assert record.openhound_version == "4.5.6" assert record.job_id == 42 - - assert all( - "traversable" not in properties - for properties in mock_bloodhound_api.app.state.ingested_edge_properties - )