MAINT: Fixing low-hanging test warnings (microsoft#1160)

Author: rlundeen2

pyrit/prompt_converter/transparency_attack_converter.py

@@ -206,7 +206,8 @@ async def _save_blended_image(self, attack_image: numpy.ndarray, alpha: numpy.nd
             img_serializer.file_extension = "png"
 
             la_image = self._create_blended_image(attack_image, alpha)
-            la_pil = Image.fromarray(la_image, mode="LA")
+            # Pillow 10+ can infer the mode from the array dtype and shape
+            la_pil = Image.fromarray(la_image)
             image_buffer = BytesIO()
             la_pil.save(image_buffer, format="PNG")
             image_str = base64.b64encode(image_buffer.getvalue())

tests/unit/converter/test_math_prompt_converter.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -12,8 +12,9 @@
 
 @pytest.mark.asyncio
 async def test_math_prompt_converter_convert_async():
-    # Mock the converter target
-    mock_converter_target = AsyncMock()
+    # Mock the converter target - use MagicMock for synchronous methods
+    mock_converter_target = MagicMock()
+    mock_converter_target.send_prompt_async = AsyncMock()
     # Specify parameters=['prompt'] to match the placeholder in the template
     template_value = "Solve the following problem: {{ prompt }}"
     dataset_name = "dataset_1"
@@ -66,8 +67,9 @@ async def test_math_prompt_converter_convert_async():
 
 @pytest.mark.asyncio
 async def test_math_prompt_converter_handles_disallowed_content():
-    # Mock the converter target
-    mock_converter_target = AsyncMock()
+    # Mock the converter target - use MagicMock for synchronous methods
+    mock_converter_target = MagicMock()
+    mock_converter_target.send_prompt_async = AsyncMock()
     # Specify parameters=['prompt'] to match the placeholder in the template
     template_value = "Encode this instruction: {{ prompt }}"
     dataset_name = "dataset_1"
@@ -117,8 +119,9 @@ async def test_math_prompt_converter_handles_disallowed_content():
 
 @pytest.mark.asyncio
 async def test_math_prompt_converter_invalid_input_type():
-    # Mock the converter target
-    mock_converter_target = AsyncMock()
+    # Mock the converter target - use MagicMock for synchronous methods
+    mock_converter_target = MagicMock()
+    mock_converter_target.send_prompt_async = AsyncMock()
     # Specify parameters=['prompt'] to match the placeholder in the template
     template_value = "Encode this instruction: {{ prompt }}"
     dataset_name = "dataset_1"
@@ -132,13 +135,15 @@ async def test_math_prompt_converter_invalid_input_type():
 
     # Test with an invalid input type
     with pytest.raises(ValueError, match="Input type not supported"):
-        await converter.convert_async(prompt="Test prompt", input_type="unsupported")
+        # Use type: ignore to suppress the type error for testing invalid input
+        await converter.convert_async(prompt="Test prompt", input_type="unsupported")  # type: ignore
 
 
 @pytest.mark.asyncio
 async def test_math_prompt_converter_error_handling():
-    # Mock the converter target
-    mock_converter_target = AsyncMock()
+    # Mock the converter target - use MagicMock for synchronous methods
+    mock_converter_target = MagicMock()
+    mock_converter_target.send_prompt_async = AsyncMock()
     # Specify parameters=['prompt'] to match the placeholder in the template
     template_value = "Encode this instruction: {{ prompt }}"
     dataset_name = "dataset_1"

tests/unit/score/test_self_ask_category.py

@@ -199,7 +199,8 @@ async def test_score_prompts_batch_async(
     scorer_category_response_false: Message,
     patch_central_database,
 ):
-    chat_target = AsyncMock()
+    chat_target = MagicMock()
+    chat_target.send_prompt_async = AsyncMock()
     chat_target._max_requests_per_minute = max_requests_per_minute
     with patch.object(CentralMemory, "get_memory_instance", return_value=MagicMock()):
         scorer = SelfAskCategoryScorer(

tests/unit/target/test_huggingface_chat_target.py

@@ -78,15 +78,23 @@ def mock_pretrained_config():
         yield
 
 
-class AwaitableMock(AsyncMock):
+class AwaitableTask(AsyncMock):
+    """Mock that can be awaited and acts like an asyncio.Task"""
+
     def __await__(self):
-        return iter([])
+        # Return a completed future-like object
+        async def _await():
+            return None
+
+        return _await().__await__()
 
 
 @pytest.fixture(autouse=True)
 def mock_create_task():
-    with patch("asyncio.create_task", return_value=AwaitableMock(spec=Task)):
-        yield
+    with patch("asyncio.create_task") as mock_task:
+        # Return an AwaitableTask that can be awaited
+        mock_task.return_value = AwaitableTask(spec=Task)
+        yield mock_task
 
 
 @pytest.fixture(autouse=True)
@@ -124,17 +132,23 @@ async def test_hf_initialization(patch_central_database, mock_download_specific_
 
 
 @pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed")
-def test_is_model_id_valid_true():
+@pytest.mark.asyncio
+async def test_is_model_id_valid_true():
     # Simulate valid model ID
     hf_chat = HuggingFaceChatTarget(model_id="test_model", use_cuda=False)
+    # Await the background task to prevent warnings
+    await hf_chat.load_model_and_tokenizer_task
     assert hf_chat.is_model_id_valid()
 
 
 @pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed")
-def test_is_model_id_valid_false():
+@pytest.mark.asyncio
+async def test_is_model_id_valid_false():
     # Simulate invalid model ID by causing an exception
     with patch("transformers.PretrainedConfig.from_pretrained", side_effect=Exception("Invalid model")):
         hf_chat = HuggingFaceChatTarget(model_id="test_model", use_cuda=False)
+        # Await the background task to prevent warnings
+        await hf_chat.load_model_and_tokenizer_task
         assert not hf_chat.is_model_id_valid()
 
 
@@ -191,8 +205,11 @@ async def test_missing_chat_template_error():
 
 
 @pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed")
-def test_invalid_prompt_request_validation():
+@pytest.mark.asyncio
+async def test_invalid_prompt_request_validation():
     hf_chat = HuggingFaceChatTarget(model_id="test_model", use_cuda=False)
+    # Await the background task to prevent warnings
+    await hf_chat.load_model_and_tokenizer_task
 
     # Create an invalid prompt request with multiple message pieces
     message_piece1 = MessagePiece(
@@ -301,18 +318,24 @@ async def test_optional_kwargs_args_passed_when_loading_model(mock_transformers)
 
 
 @pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed")
-def test_is_json_response_supported():
+@pytest.mark.asyncio
+async def test_is_json_response_supported():
     hf_chat = HuggingFaceChatTarget(model_id="dummy", use_cuda=False, trust_remote_code=True)
+    # Await the background task to prevent warnings
+    await hf_chat.load_model_and_tokenizer_task
     assert hf_chat.is_json_response_supported() is False
 
 
 @pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed")
-def test_hugging_face_chat_sets_endpoint_and_rate_limit():
+@pytest.mark.asyncio
+async def test_hugging_face_chat_sets_endpoint_and_rate_limit():
     target = HuggingFaceChatTarget(
         model_id="test_model",
         use_cuda=False,
         max_requests_per_minute=30,
     )
+    # Await the background task to prevent warnings
+    await target.load_model_and_tokenizer_task
     identifier = target.get_identifier()
     # HuggingFaceChatTarget doesn't set an endpoint (it's local), so it shouldn't be in identifier
     assert "endpoint" not in identifier

tests/unit/test_azure_storage_auth.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 from datetime import datetime, timedelta
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from azure.storage.blob import UserDelegationKey
@@ -54,7 +54,8 @@ async def test_get_sas_token(
 
     # Mocking the container URL
     container_url = "https://mockaccount.blob.core.windows.net/mockcontainer"
-    mock_container_client_instance = AsyncMock()
+    # Use MagicMock with spec to avoid AsyncMock behavior for sync properties
+    mock_container_client_instance = MagicMock(spec=["container_name", "account_name"])
     mock_container_client.from_container_url.return_value = mock_container_client_instance
     mock_container_client_instance.container_name = "mock_container"
     mock_container_client_instance.account_name = "mock_account"

tests/unit/test_hf_model_downloads.py

@@ -32,10 +32,11 @@ def setup_environment():
         yield token
 
 
-def test_download_specific_files(setup_environment):
+@pytest.mark.asyncio
+async def test_download_specific_files(setup_environment):
     """Test downloading specific files"""
     token = setup_environment  # Get the token from the fixture
 
     with patch("os.makedirs"):
         with patch("pyrit.common.download_hf_model.download_files"):
-            download_specific_files(MODEL_ID, FILE_PATTERNS, token, Path(""))
+            await download_specific_files(MODEL_ID, FILE_PATTERNS, token, Path(""))

tests/unit/test_prompt_normalizer.py

@@ -124,7 +124,8 @@ async def test_send_prompt_async_no_response_adds_memory(mock_memory_instance, s
 
 @pytest.mark.asyncio
 async def test_send_prompt_async_empty_response_exception_handled(mock_memory_instance, seed_group):
-    prompt_target = AsyncMock()
+    # Use MagicMock with send_prompt_async as AsyncMock to avoid coroutine warnings on other methods
+    prompt_target = MagicMock()
     prompt_target.send_prompt_async = AsyncMock(side_effect=EmptyResponseException(message="Empty response"))
 
     normalizer = PromptNormalizer()
@@ -142,7 +143,8 @@ async def test_send_prompt_async_empty_response_exception_handled(mock_memory_in
 
 @pytest.mark.asyncio
 async def test_send_prompt_async_request_response_added_to_memory(mock_memory_instance, seed_group):
-    prompt_target = AsyncMock()
+    # Use MagicMock with send_prompt_async as AsyncMock to avoid coroutine warnings
+    prompt_target = MagicMock()
 
     response = MessagePiece(role="assistant", original_value="test_response").to_message()
 
@@ -288,6 +290,9 @@ async def test_send_prompt_async_converters_response(mock_memory_instance, seed_
 @pytest.mark.asyncio
 async def test_send_prompt_async_image_converter(mock_memory_instance):
     prompt_target = MagicMock(PromptTarget)
+    prompt_target.send_prompt_async = AsyncMock(
+        return_value=MessagePiece(role="assistant", original_value="response").to_message()
+    )
 
     mock_image_converter = MagicMock(PromptConverter)
 
@@ -297,9 +302,11 @@ async def test_send_prompt_async_image_converter(mock_memory_instance):
         filename = f.name
         f.write(b"Hello")
 
-        mock_image_converter.convert_tokens_async.return_value = ConverterResult(
-            output_type="image_path",
-            output_text=filename,
+        mock_image_converter.convert_tokens_async = AsyncMock(
+            return_value=ConverterResult(
+                output_type="image_path",
+                output_text=filename,
+            )
         )
 
         prompt_converters = PromptConverterConfiguration(converters=[mock_image_converter])