From d8a7a7198e6c745725e3ea085af5b14f815a27bb Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Thu, 12 Mar 2026 19:59:42 -0700 Subject: [PATCH] Fix type annotation warnings and test warnings (issue #442) - Add -> None return type to TranslationConverter.__init__ - Add Optional[str] annotation to AzureBlobStorageTarget._sas_token - Fix _client_async type to Optional[AsyncContainerClient] - Replace unclosed open().read() with Path.read_bytes() in test_tts_target - Use context manager for Image.open() in test_leakage_scenario - Replace deprecated .role with .api_role in test_websocket_copilot_target - Use 32+ byte JWT key to fix InsecureKeyLengthWarning in tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/prompt_converter/translation_converter.py | 2 +- pyrit/prompt_target/azure_blob_storage_target.py | 4 ++-- tests/unit/scenarios/test_leakage_scenario.py | 10 +++++----- tests/unit/target/test_tts_target.py | 3 ++- tests/unit/target/test_websocket_copilot_target.py | 8 ++++---- 5 files changed, 14 insertions(+), 13 deletions(-) diff --git a/pyrit/prompt_converter/translation_converter.py b/pyrit/prompt_converter/translation_converter.py index 7b8c4c7c49..911f72ab57 100644 --- a/pyrit/prompt_converter/translation_converter.py +++ b/pyrit/prompt_converter/translation_converter.py @@ -46,7 +46,7 @@ def __init__( prompt_template: Optional[SeedPrompt] = None, max_retries: int = 3, max_wait_time_in_seconds: int = 60, - ): + ) -> None: """ Initialize the converter with the target chat support, language, and optional prompt template. diff --git a/pyrit/prompt_target/azure_blob_storage_target.py b/pyrit/prompt_target/azure_blob_storage_target.py index 7303865c10..824c104f47 100644 --- a/pyrit/prompt_target/azure_blob_storage_target.py +++ b/pyrit/prompt_target/azure_blob_storage_target.py @@ -75,8 +75,8 @@ def __init__( env_var_name=self.AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE, passed_value=container_url ) - self._sas_token = sas_token - self._client_async: AsyncContainerClient = None + self._sas_token: Optional[str] = sas_token + self._client_async: Optional[AsyncContainerClient] = None super().__init__(endpoint=self._container_url, max_requests_per_minute=max_requests_per_minute) diff --git a/tests/unit/scenarios/test_leakage_scenario.py b/tests/unit/scenarios/test_leakage_scenario.py index 1641af0d23..b7b7d066db 100644 --- a/tests/unit/scenarios/test_leakage_scenario.py +++ b/tests/unit/scenarios/test_leakage_scenario.py @@ -536,9 +536,9 @@ def test_ensure_blank_image_exists_creates_image(self, mock_objective_scorer, sa # Verify it's a valid image with correct dimensions from PIL import Image - img = Image.open(test_image_path) - assert img.size == (800, 600) - assert img.mode == "RGB" + with Image.open(test_image_path) as img: + assert img.size == (800, 600) + assert img.mode == "RGB" def test_ensure_blank_image_exists_does_not_overwrite(self, mock_objective_scorer, sample_objectives, tmp_path): """Test that _ensure_blank_image_exists doesn't overwrite existing image.""" @@ -564,8 +564,8 @@ def test_ensure_blank_image_exists_does_not_overwrite(self, mock_objective_score assert Path(test_image_path).stat().st_mtime == original_mtime # Verify it's still the original image - img = Image.open(test_image_path) - assert img.size == (100, 100) # Original size, not 800x600 + with Image.open(test_image_path) as img: + assert img.size == (100, 100) # Original size, not 800x600 def test_ensure_blank_image_exists_creates_parent_directories( self, mock_objective_scorer, sample_objectives, tmp_path diff --git a/tests/unit/target/test_tts_target.py b/tests/unit/target/test_tts_target.py index 70b53dd6c6..42ccc8a1b0 100644 --- a/tests/unit/target/test_tts_target.py +++ b/tests/unit/target/test_tts_target.py @@ -4,6 +4,7 @@ import os import uuid from collections.abc import MutableSequence +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -121,7 +122,7 @@ async def test_tts_send_prompt_file_save_async( assert file_path assert file_path.endswith(f".{response_format}") assert os.path.exists(file_path) - data = open(file_path, "rb").read() # noqa: SIM115 + data = Path(file_path).read_bytes() assert data == b"audio data" os.remove(file_path) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index 21d3e610b7..e44f4e73d3 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -16,7 +16,7 @@ @pytest.fixture def mock_authenticator(): token_payload = {"tid": "test_tenant_id", "oid": "test_object_id", "exp": 9999999999} - mock_token = jwt.encode(token_payload, "secret", algorithm="HS256") + mock_token = jwt.encode(token_payload, "a]T3!s$ecReT#kEy_4_tEst!ng@2024+", algorithm="HS256") if isinstance(mock_token, bytes): mock_token = mock_token.decode("utf-8") authenticator = MagicMock(spec=CopilotAuthenticator) @@ -267,7 +267,7 @@ async def test_build_websocket_url_with_missing_ids(self, mock_authenticator, mo token_payload = {"tid": "test_tenant_id", "oid": "test_object_id", "exp": 9999999999} del token_payload[missing_id] - mock_token = jwt.encode(token_payload, "secret", algorithm="HS256") + mock_token = jwt.encode(token_payload, "a]T3!s$ecReT#kEy_4_tEst!ng@2024+", algorithm="HS256") if isinstance(mock_token, bytes): mock_token = mock_token.decode("utf-8") mock_authenticator.get_token = AsyncMock(return_value=mock_token) @@ -848,7 +848,7 @@ async def test_send_prompt_async_successful(self, mock_authenticator, make_messa assert len(responses) == 1 assert responses[0].message_pieces[0].converted_value == "Response from Copilot" - assert responses[0].message_pieces[0].role == "assistant" + assert responses[0].message_pieces[0].api_role == "assistant" @pytest.mark.asyncio async def test_send_prompt_async_with_exceptions(self, mock_authenticator, make_message_piece, mock_memory): @@ -884,7 +884,7 @@ async def test_send_prompt_async_with_image(self, mock_authenticator, make_messa assert len(responses) == 1 assert responses[0].message_pieces[0].converted_value == "Image description response" - assert responses[0].message_pieces[0].role == "assistant" + assert responses[0].message_pieces[0].api_role == "assistant" @pytest.mark.asyncio async def test_send_prompt_async_with_mixed_content(self, mock_authenticator, make_message_piece, mock_memory):