diff --git a/lib/group_repository.py b/lib/group_repository.py index 441b1e9171..df426f96b1 100644 --- a/lib/group_repository.py +++ b/lib/group_repository.py @@ -576,27 +576,25 @@ def get_ungrouped_group(identity: Identity) -> Group: return ungrouped_group -def serialize_group(group: Group | dict, org_id: str, account: str | None = None) -> dict: +def serialize_group( + group: Group | dict, org_id: str, account: str | None = None, *, with_host_count: bool = True +) -> dict: """ - Serialize a group with host count. + Serialize a group, optionally including host count. Delegates to the appropriate serializer based on whether the group is from the database or RBAC v2. Args: group: Either a Group ORM object (from DB) or a dict (from RBAC v2) org_id: The organization ID account: The account_number (optional, only used for RBAC v2 workspaces) - - Returns: - Dictionary containing serialized group data with host_count + with_host_count: When False, skip the host count query (used during MQ ingestion + where an accurate count is not needed and the query is expensive). """ if isinstance(group, dict): - # RBAC v2 workspace (dict from RBAC API) - # Extract group_id from dict and get host count using batch function group_id = group["id"] - host_count = get_host_counts_batch(org_id, [group_id])[group_id] + host_count = get_host_counts_batch(org_id, [group_id])[group_id] if with_host_count else 0 return serialize_rbac_workspace_with_host_count(group, org_id, account, host_count) else: - # Database Group (ORM object) group_id = str(group.id) - host_count = get_host_counts_batch(org_id, [group_id])[group_id] + host_count = get_host_counts_batch(org_id, [group_id])[group_id] if with_host_count else 0 return serialize_db_group_with_host_count(group, host_count) diff --git a/lib/host_repository.py b/lib/host_repository.py index 62fa360035..9dbffce9f8 100644 --- a/lib/host_repository.py +++ b/lib/host_repository.py @@ -90,7 +90,9 @@ def add_host( matched_host = find_existing_host(identity, canonical_facts) group = get_or_create_ungrouped_hosts_group_for_identity(identity) - input_host.groups = [serialize_group(group, identity.org_id, getattr(identity, "account_number", None))] + input_host.groups = [ + serialize_group(group, identity.org_id, getattr(identity, "account_number", None), with_host_count=False) + ] if matched_host: defer_to_reporter = operation_args.get("defer_to_reporter") diff --git a/tests/test_models.py b/tests/test_models.py index 65d76143d4..42becdb71b 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -2719,3 +2719,46 @@ def test_update_display_name_writes_when_reporter_changed(db_create_host, models assert existing_host.display_name == "my-host" assert existing_host.display_name_reporter == "yupana" + + +def test_serialize_group_with_host_count_false(flask_app, mocker): # noqa: ARG001 + """serialize_group(with_host_count=False) should skip the host count query.""" + from lib.group_repository import serialize_group + + mock_group = mocker.Mock() + mock_group.id = "test-group-id" + mock_group.name = "Test Group" + mock_group.org_id = "test_org" + + count_spy = mocker.patch("lib.group_repository.get_host_counts_batch") + mocker.patch( + "lib.group_repository.serialize_db_group_with_host_count", + return_value={"id": "test-group-id", "name": "Test Group", "host_count": 0}, + ) + + result = serialize_group(mock_group, "test_org", with_host_count=False) + + count_spy.assert_not_called() + assert result["host_count"] == 0 + + +def test_serialize_group_with_host_count_true(flask_app, mocker): # noqa: ARG001 + """serialize_group(with_host_count=True) should query for host count (default behavior).""" + from lib.group_repository import serialize_group + + mock_group = mocker.Mock() + mock_group.id = "test-group-id" + + count_spy = mocker.patch( + "lib.group_repository.get_host_counts_batch", + return_value={"test-group-id": 42}, + ) + mocker.patch( + "lib.group_repository.serialize_db_group_with_host_count", + return_value={"id": "test-group-id", "host_count": 42}, + ) + + result = serialize_group(mock_group, "test_org") + + count_spy.assert_called_once() + assert result["host_count"] == 42