diff --git a/pyrit/analytics/attack_result_analytics.py b/pyrit/analytics/attack_result_analytics.py index df15cdff61..96200fded8 100644 --- a/pyrit/analytics/attack_result_analytics.py +++ b/pyrit/analytics/attack_result_analytics.py @@ -6,7 +6,7 @@ from __future__ import annotations import threading -from typing import TYPE_CHECKING, ClassVar, Self, TypeVar +from typing import TYPE_CHECKING, ClassVar, Self, TypeVar, assert_never from weakref import WeakKeyDictionary from pydantic import BaseModel @@ -305,20 +305,33 @@ def _option(*, raw: RawAnalyticsOption, dimension: AttackAnalyticsDimension) -> label = "Not recorded" elif raw.key.kind is AttackAnalyticsValueKind.NO_CONVERTERS: label = "No converters" + elif raw.key.kind is AttackAnalyticsValueKind.VALUE: + label = AttackResultAnalytics._value_label(raw=raw, dimension=dimension) + elif raw.key.kind is AttackAnalyticsValueKind.LEGACY_OPERATION_NAME: + label = f"{AttackResultAnalytics._value_label(raw=raw, dimension=dimension)} (legacy name)" else: - label = raw.label if raw.label is not None else raw.key.value or "" - if not label.strip(): - label = "(Blank)" - elif ( - dimension.name in {AttackAnalyticsDimensionName.OBJECTIVE_TARGET, AttackAnalyticsDimensionName.SCENARIO} - and raw.key.value is not None - and label != raw.key.value - ): - label = f"{label} ({raw.key.value[-8:]})" - if raw.key.kind is AttackAnalyticsValueKind.LEGACY_OPERATION_NAME: - label = f"{label} (legacy name)" + assert_never(raw.key.kind) return AttackAnalyticsOption(key=raw.key, label=label) + @staticmethod + def _value_label(*, raw: RawAnalyticsOption, dimension: AttackAnalyticsDimension) -> str: + """ + Format the text-bearing value kinds without changing their stored keys. + + Returns: + str: The SDK display label for a text-bearing value kind. + """ + label = raw.label if raw.label is not None else raw.key.value or "" + if not label.strip(): + return "(Blank)" + if ( + dimension.name in {AttackAnalyticsDimensionName.OBJECTIVE_TARGET, AttackAnalyticsDimensionName.SCENARIO} + and raw.key.value is not None + and label != raw.key.value + ): + return f"{label} ({raw.key.value[-8:]})" + return label + @staticmethod def _drilldown_unavailable_reason(query: AttackAnalyticsQuery) -> str | None: """ diff --git a/tests/unit/analytics/test_attack_result_analytics.py b/tests/unit/analytics/test_attack_result_analytics.py index f8d922be7c..f18d79b8d9 100644 --- a/tests/unit/analytics/test_attack_result_analytics.py +++ b/tests/unit/analytics/test_attack_result_analytics.py @@ -4,7 +4,12 @@ from __future__ import annotations import asyncio +import inspect +import subprocess +import sys from datetime import UTC, datetime, timedelta, timezone +from pathlib import Path +from textwrap import dedent from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -28,6 +33,7 @@ AttackAnalyticsQuery, AttackAnalyticsResultsQuery, AttackAnalyticsValue, + AttackAnalyticsValueKind, AttackIdentifier, AttackOutcome, AttackResultMetadata, @@ -327,6 +333,7 @@ async def test_matrix_empty_cells_and_additional_predicates_are_exact( ("value", " \t", None, "label", "(Blank)"), ("value", "Unknown", "Unknown", "operation", "Unknown"), ("value", "Not recorded", "Not recorded", "operation", "Not recorded"), + ("legacy_operation_name", "Legacy", "Legacy", "operation", "Legacy (legacy name)"), ("value", "content-independent-eval-hash", "model", "objective_target", "model (val-hash)"), ("value", "scenario-12345678", "run", "scenario", "run (12345678)"), ("value", "scenario-12345678", None, "scenario", "scenario-12345678"), @@ -344,6 +351,71 @@ def test_display_labels_do_not_change_typed_keys( assert option.key == raw.key +@pytest.mark.parametrize("extend_enum", [False, True], ids=["current-enum", "extended-enum"]) +def test_option_dispatch_type_exhaustiveness(tmp_path: Path, extend_enum: bool) -> None: + enum_source = inspect.getsource(AttackAnalyticsValueKind) + if extend_enum: + enum_source += '\n FUTURE_KIND = "future_kind"\n' + probe = tmp_path / "dispatch.py" + probe.write_text( + "\n".join( + [ + "from dataclasses import dataclass", + "from enum import Enum", + "from typing import assert_never", + "from pyrit.models import AttackAnalyticsDimension, AttackAnalyticsDimensionName", + enum_source, + dedent( + """ + @dataclass + class AttackAnalyticsValue: + kind: AttackAnalyticsValueKind + value: str | None + + @dataclass + class RawAnalyticsOption: + key: AttackAnalyticsValue + label: str | None + + @dataclass + class AttackAnalyticsOption: + key: AttackAnalyticsValue + label: str + """ + ), + "class AttackResultAnalytics:", + inspect.getsource(AttackResultAnalytics._option), + inspect.getsource(AttackResultAnalytics._value_label), + ] + ), + encoding="utf-8", + ) + result = subprocess.run( + [ + sys.executable, + "-m", + "ty", + "check", + "--project", + str(Path(__file__).resolve().parents[3]), + "--python-version", + "3.11", + str(probe), + ], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + output = result.stdout + result.stderr + assert result.returncode == (1 if extend_enum else 0), output + if extend_enum: + assert output.count("error[") == 1, output + assert "type-assertion-failure" in output, output + assert "assert_never(raw.key.kind)" in output, output + assert "FUTURE_KIND" in output, output + + async def test_operation_option_labels(analytics: AttackResultAnalytics, sqlite_instance: SQLiteMemory) -> None: saved, removed = [await sqlite_instance.add_operation_async(Operation(name=name)) for name in ("Saved", "Gone")] await add_legacy_attack_results_async(