Fix lint errors in test_custom_llm.py

- Add noqa comment for hardcoded test JWT token
- Add return statement to satisfy ruff RET503 check

Co-Authored-By: João <joao@crewai.com>

tests/test_custom_llm.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import pytest
 
@@ -159,11 +159,11 @@ class JWTAuthLLM(BaseLLM):
 
     def call(
         self,
-        messages: Union[str, List[Dict[str, str]]],
+        messages: str | list[dict[str, str]],
-        tools: Optional[List[dict]] = None,
+        tools: list[dict] | None = None,
-        callbacks: Optional[List[Any]] = None,
+        callbacks: list[Any] | None = None,
-        available_functions: Optional[Dict[str, Any]] = None,
+        available_functions: dict[str, Any] | None = None,
-    ) -> Union[str, Any]:
+    ) -> str | Any:
         """Record the call and return a predefined response."""
         self.calls.append(
             {
@@ -192,7 +192,7 @@ class JWTAuthLLM(BaseLLM):
 
 def test_custom_llm_with_jwt_auth():
     """Test a custom LLM implementation with JWT authentication."""
-    jwt_llm = JWTAuthLLM(jwt_token="example.jwt.token")
+    jwt_llm = JWTAuthLLM(jwt_token="example.jwt.token")  # noqa: S106
 
     # Test that create_llm returns the JWT-authenticated LLM instance directly
     result_llm = create_llm(jwt_llm)
@@ -238,11 +238,11 @@ class TimeoutHandlingLLM(BaseLLM):
 
     def call(
         self,
-        messages: Union[str, List[Dict[str, str]]],
+        messages: str | list[dict[str, str]],
-        tools: Optional[List[dict]] = None,
+        tools: list[dict] | None = None,
-        callbacks: Optional[List[Any]] = None,
+        callbacks: list[Any] | None = None,
-        available_functions: Optional[Dict[str, Any]] = None,
+        available_functions: dict[str, Any] | None = None,
-    ) -> Union[str, Any]:
+    ) -> str | Any:
         """Simulate API calls with timeout handling and retry logic.
 
         Args:
@@ -282,35 +282,34 @@ class TimeoutHandlingLLM(BaseLLM):
                         )
                     # Otherwise, continue to the next attempt (simulating backoff)
                     continue
-                else:
+                # Success on first attempt
-                    # Success on first attempt
+                return "First attempt response"
-                    return "First attempt response"
+            # This is a retry attempt (attempt > 0)
-            else:
+            # Always record retry attempts
-                # This is a retry attempt (attempt > 0)
+            self.calls.append(
-                # Always record retry attempts
+                {
-                self.calls.append(
+                    "retry_attempt": attempt,
-                    {
+                    "messages": messages,
-                        "retry_attempt": attempt,
+                    "tools": tools,
-                        "messages": messages,
+                    "callbacks": callbacks,
-                        "tools": tools,
+                    "available_functions": available_functions,
-                        "callbacks": callbacks,
+                }
-                        "available_functions": available_functions,
+            )
-                    }
-                )
 
-                # Simulate a failure if fail_count > 0
+            # Simulate a failure if fail_count > 0
-                if self.fail_count > 0:
+            if self.fail_count > 0:
-                    self.fail_count -= 1
+                self.fail_count -= 1
-                    # If we've used all retries, raise an error
+                # If we've used all retries, raise an error
-                    if attempt == self.max_retries - 1:
+                if attempt == self.max_retries - 1:
-                        raise TimeoutError(
+                    raise TimeoutError(
-                            f"LLM request failed after {self.max_retries} attempts"
+                        f"LLM request failed after {self.max_retries} attempts"
-                        )
+                    )
-                    # Otherwise, continue to the next attempt (simulating backoff)
+                # Otherwise, continue to the next attempt (simulating backoff)
-                    continue
+                continue
-                else:
+            # Success on retry
-                    # Success on retry
+            return "Response after retry"
-                    return "Response after retry"
+
+        return "Response after retry"
 
     def supports_function_calling(self) -> bool:
         """Return True to indicate that function calling is supported.
@@ -368,11 +367,11 @@ class MinimalCustomLLM(BaseLLM):
 
     def call(
         self,
-        messages: Union[str, List[Dict[str, str]]],
+        messages: str | list[dict[str, str]],
-        tools: Optional[List[dict]] = None,
+        tools: list[dict] | None = None,
-        callbacks: Optional[List[Any]] = None,
+        callbacks: list[Any] | None = None,
-        available_functions: Optional[Dict[str, Any]] = None,
+        available_functions: dict[str, Any] | None = None,
-    ) -> Union[str, Any]:
+    ) -> str | Any:
         return "Minimal response"