Style fixing for tests/ (#8225)

* enable style check

* update rules

* allow wildcard imports

* fix tests style

Author: chenmoneygithub

tests/adapters/test_json_adapter.py

@@ -2,8 +2,7 @@
 
 import pydantic
 import pytest
-from pydantic import create_model
-from litellm.utils import ModelResponse, Message, Choices
+from litellm.utils import Choices, Message, ModelResponse
 
 import dspy
 
@@ -109,7 +108,7 @@ async def test_json_adapter_async_call():
 
 
 def test_json_adapter_on_pydantic_model():
-    from litellm.utils import ModelResponse, Message, Choices
+    from litellm.utils import Choices, Message, ModelResponse
 
     class User(pydantic.BaseModel):
         id: int
@@ -170,7 +169,7 @@ class TestSignature(dspy.Signature):
         assert expected_input_structure in content
 
         # Assert that system prompt includes output formatting structure
-        expected_output_structure = (  # noqa: Q000
+        expected_output_structure = (
             "Outputs will be a JSON object with the following fields.\n\n{\n  "
             '"answer": "{answer}        # note: the value you produce must adhere to the JSON schema: '
             '{\\"type\\": \\"object\\", \\"properties\\": {\\"analysis\\": {\\"type\\": \\"string\\", \\"title\\": '
@@ -184,7 +183,7 @@ class TestSignature(dspy.Signature):
         assert user_message_content is not None
 
         # Assert that the user input data is formatted correctly
-        expected_input_data = (  # noqa: Q000
+        expected_input_data = (
             '[[ ## user ## ]]\n{"id": 5, "name": "name_test", "email": "email_test"}\n\n[[ ## question ## ]]\n'
             "What is the capital of France?\n\n"
         )

tests/adapters/test_two_step_adapter.py

@@ -1,4 +1,5 @@
 from unittest import mock
+
 import pytest
 
 import dspy
@@ -9,7 +10,7 @@ class TestSignature(dspy.Signature):
         question: str = dspy.InputField(desc="The math question to solve")
         solution: str = dspy.OutputField(desc="Step by step solution")
         answer: float = dspy.OutputField(desc="The final numerical answer")
-    
+
     program = dspy.Predict(TestSignature)
 
     mock_main_lm = mock.MagicMock(spec=dspy.LM)
@@ -30,7 +31,7 @@ class TestSignature(dspy.Signature):
     dspy.configure(lm=mock_main_lm, adapter=dspy.TwoStepAdapter(extraction_model=mock_extraction_lm))
 
     result = program(question="What is 5 + 7?")
-    
+
     assert result.answer == 12
 
     # main LM call
@@ -59,22 +60,23 @@ class TestSignature(dspy.Signature):
     # assert first message
     assert call_kwargs["messages"][0]["role"] == "system"
     content = call_kwargs["messages"][0]["content"]
-    assert  "`text` (str)" in content
-    assert  "`solution` (str)" in content
-    assert  "`answer` (float)" in content
+    assert "`text` (str)" in content
+    assert "`solution` (str)" in content
+    assert "`answer` (float)" in content
 
     # assert second message
     assert call_kwargs["messages"][1]["role"] == "user"
     content = call_kwargs["messages"][1]["content"]
     assert "text from main LM" in content
 
+
 @pytest.mark.asyncio
 async def test_two_step_adapter_async_call():
     class TestSignature(dspy.Signature):
         question: str = dspy.InputField(desc="The math question to solve")
         solution: str = dspy.OutputField(desc="Step by step solution")
         answer: float = dspy.OutputField(desc="The final numerical answer")
-    
+
     program = dspy.Predict(TestSignature)
 
     mock_main_lm = mock.MagicMock(spec=dspy.LM)
@@ -95,7 +97,7 @@ class TestSignature(dspy.Signature):
     dspy.configure(lm=mock_main_lm, adapter=dspy.TwoStepAdapter(extraction_model=mock_extraction_lm))
 
     result = await program.acall(question="What is 5 + 7?")
-    
+
     assert result.answer == 12
 
     # main LM call
@@ -124,37 +126,40 @@ class TestSignature(dspy.Signature):
     # assert first message
     assert call_kwargs["messages"][0]["role"] == "system"
     content = call_kwargs["messages"][0]["content"]
-    assert  "`text` (str)" in content
-    assert  "`solution` (str)" in content
-    assert  "`answer` (float)" in content
+    assert "`text` (str)" in content
+    assert "`solution` (str)" in content
+    assert "`answer` (float)" in content
 
     # assert second message
     assert call_kwargs["messages"][1]["role"] == "user"
     content = call_kwargs["messages"][1]["content"]
     assert "text from main LM" in content
 
+
 def test_two_step_adapter_parse():
     class ComplexSignature(dspy.Signature):
         input_text: str = dspy.InputField()
         tags: list[str] = dspy.OutputField(desc="List of relevant tags")
         confidence: float = dspy.OutputField(desc="Confidence score")
-    
+
     first_response = "main LM response"
-    
+
     mock_extraction_lm = mock.MagicMock(spec=dspy.LM)
-    mock_extraction_lm.return_value = ["""
+    mock_extraction_lm.return_value = [
+        """
         {
             "tags": ["AI", "deep learning", "neural networks"],
             "confidence": 0.87
-        }   
-    """]
+        }
+    """
+    ]
     mock_extraction_lm.kwargs = {"temperature": 1.0}
     mock_extraction_lm.model = "openai/gpt-4o"
     adapter = dspy.TwoStepAdapter(mock_extraction_lm)
     dspy.configure(adapter=adapter, lm=mock_extraction_lm)
 
     result = adapter.parse(ComplexSignature, first_response)
-    
+
     assert result["tags"] == ["AI", "deep learning", "neural networks"]
     assert result["confidence"] == 0.87
 
@@ -163,7 +168,7 @@ def test_two_step_adapter_parse_errors():
     class TestSignature(dspy.Signature):
         question: str = dspy.InputField()
         answer: str = dspy.OutputField()
-    
+
     first_response = "main LM response"
 
     mock_extraction_lm = mock.MagicMock(spec=dspy.LM)
@@ -172,6 +177,6 @@ class TestSignature(dspy.Signature):
     mock_extraction_lm.model = "openai/gpt-4o"
 
     adapter = dspy.TwoStepAdapter(mock_extraction_lm)
-    
+
     with pytest.raises(ValueError, match="Failed to parse response"):
-        adapter.parse(TestSignature, first_response) 
+        adapter.parse(TestSignature, first_response)

tests/caching/test_litellm_cache.py

@@ -7,7 +7,7 @@
 import pytest
 
 import dspy
-from tests.test_utils.server import litellm_test_server, read_litellm_test_server_request_logs
+from tests.test_utils.server import read_litellm_test_server_request_logs
 
 
 @pytest.fixture()
@@ -158,8 +158,8 @@ class NonJsonSerializable:
 
 
 def test_lms_called_expected_number_of_times_for_cache_key_generation_failures():
-    with pytest.raises(Exception), patch("litellm.completion") as mock_completion:
-        mock_completion.side_effect = Exception("Mocked exception")
+    with pytest.raises(RuntimeError), patch("litellm.completion") as mock_completion:
+        mock_completion.side_effect = RuntimeError("Mocked exception")
         lm = dspy.LM(
             model="openai/dspy-test-model",
             api_base="fakebase",

tests/clients/test_cache.py

@@ -1,5 +1,4 @@
-import os
-import shutil
+from dataclasses import dataclass
 from unittest.mock import patch
 
 import pydantic
@@ -8,7 +7,6 @@
 from diskcache import FanoutCache
 
 from dspy.clients.cache import Cache
-from dataclasses import dataclass
 
 
 @dataclass

tests/clients/test_databricks.py

@@ -4,15 +4,15 @@
 manual testing.
 """
 
+import pytest
+
+import dspy
 from dspy.clients.databricks import (
     DatabricksProvider,
-    _create_directory_in_databricks_unity_catalog,
     TrainingJobDatabricks,
+    _create_directory_in_databricks_unity_catalog,
 )
 
-import pytest
-import dspy
-
 try:
     from databricks.sdk import WorkspaceClient
 
@@ -66,7 +66,7 @@ def test_create_finetuning_job():
 
     job = TrainingJobDatabricks()
 
-    finetuned_model = DatabricksProvider.finetune(
+    DatabricksProvider.finetune(
         job=job,
         model="meta-llama/Llama-3.2-1B",
         train_data=fake_training_data,

tests/clients/test_embedding.py

@@ -1,9 +1,10 @@
-import pytest
-from unittest.mock import Mock, patch
+from unittest.mock import patch
+
 import numpy as np
+import pytest
 
-from dspy.clients.embedding import Embedder
 import dspy
+from dspy.clients.embedding import Embedder
 
 
 # Mock response format similar to litellm's embedding response.

tests/clients/test_inspect_global_history.py

@@ -1,23 +1,26 @@
 import pytest
-from dspy.utils.dummies import DummyLM
-from dspy.clients.base_lm import GLOBAL_HISTORY
+
 import dspy
+from dspy.clients.base_lm import GLOBAL_HISTORY
+from dspy.utils.dummies import DummyLM
+
 
 @pytest.fixture(autouse=True)
 def clear_history():
     GLOBAL_HISTORY.clear()
     yield
 
+
 def test_inspect_history_basic(capsys):
     # Configure a DummyLM with some predefined responses
     lm = DummyLM([{"response": "Hello"}, {"response": "How are you?"}])
     dspy.settings.configure(lm=lm)
-    
+
     # Make some calls to generate history
     predictor = dspy.Predict("query: str -> response: str")
     predictor(query="Hi")
     predictor(query="What's up?")
-    
+
     # Test inspecting all history
     history = GLOBAL_HISTORY
     print(capsys)
@@ -26,45 +29,48 @@ def test_inspect_history_basic(capsys):
     assert all(isinstance(entry, dict) for entry in history)
     assert all("messages" in entry for entry in history)
 
+
 def test_inspect_history_with_n(capsys):
     """Test that inspect_history works with n
     Random failures in this test most likely mean you are printing messages somewhere
     """
     lm = DummyLM([{"response": "One"}, {"response": "Two"}, {"response": "Three"}])
     dspy.settings.configure(lm=lm)
-    
+
     # Generate some history
     predictor = dspy.Predict("query: str -> response: str")
     predictor(query="First")
     predictor(query="Second")
     predictor(query="Third")
-    
+
     dspy.inspect_history(n=2)
     # Test getting last 2 entries
     out, err = capsys.readouterr()
-    assert not "First" in out
+    assert "First" not in out
     assert "Second" in out
     assert "Third" in out
 
+
 def test_inspect_empty_history(capsys):
     # Configure fresh DummyLM
     lm = DummyLM([])
     dspy.settings.configure(lm=lm)
-    
+
     # Test inspecting empty history
     dspy.inspect_history()
     history = GLOBAL_HISTORY
     assert len(history) == 0
     assert isinstance(history, list)
 
+
 def test_inspect_history_n_larger_than_history(capsys):
     lm = DummyLM([{"response": "First"}, {"response": "Second"}])
     dspy.settings.configure(lm=lm)
-    
+
     predictor = dspy.Predict("query: str -> response: str")
     predictor(query="Query 1")
     predictor(query="Query 2")
-    
+
     # Request more entries than exist
     dspy.inspect_history(n=5)
     history = GLOBAL_HISTORY

tests/clients/test_lm.py

@@ -1,16 +1,15 @@
+import time
 from unittest import mock
 from unittest.mock import patch
 
-import time
 import litellm
 import pydantic
 import pytest
+from litellm.utils import Choices, Message, ModelResponse
 from openai import RateLimitError
 
 import dspy
 from dspy.utils.usage_tracker import track_usage
-from tests.test_utils.server import litellm_test_server, read_litellm_test_server_request_logs
-from litellm.utils import ModelResponse, Message, Choices
 
 
 def test_chat_lms_can_be_queried(litellm_test_server):
@@ -127,7 +126,10 @@ def test_lm_calls_support_callables(litellm_test_server):
     api_base, _ = litellm_test_server
 
     with mock.patch("litellm.completion", autospec=True, wraps=litellm.completion) as spy_completion:
-        azure_ad_token_provider = lambda *args, **kwargs: None
+
+        def azure_ad_token_provider(*args, **kwargs):
+            return None
+
         lm_with_callable = dspy.LM(
             model="openai/dspy-test-model",
             api_base=api_base,
@@ -315,12 +317,12 @@ def test_logprobs_included_when_requested():
                 },
             ]
         }
-        assert mock_completion.call_args.kwargs["logprobs"] == True
+        assert mock_completion.call_args.kwargs["logprobs"]
 
 
 @pytest.mark.asyncio
 async def test_async_lm_call():
-    from litellm.utils import ModelResponse, Message, Choices
+    from litellm.utils import Choices, Message, ModelResponse
 
     mock_response = ModelResponse(choices=[Choices(message=Message(content="answer"))], model="openai/gpt-4o-mini")