From cb85025abe2b7a79bd4550b8f44665dfa06c196f Mon Sep 17 00:00:00 2001 From: Jose Osorio Date: Mon, 14 Jul 2025 21:56:00 -0500 Subject: [PATCH] EPD-2096: add retry to run_test_suite --- autoblocks/_impl/testing/run.py | 55 +++- autoblocks/_impl/testing/v2/run.py | 56 +++- tests/autoblocks/test_run_test_suite.py | 276 ++++++++++++++++++ .../test_run_test_suite_v2_retry.py | 257 ++++++++++++++++ 4 files changed, 639 insertions(+), 5 deletions(-) create mode 100644 tests/autoblocks/test_run_test_suite_v2_retry.py diff --git a/autoblocks/_impl/testing/run.py b/autoblocks/_impl/testing/run.py index a8afdce0..a7c3b36e 100644 --- a/autoblocks/_impl/testing/run.py +++ b/autoblocks/_impl/testing/run.py @@ -14,6 +14,12 @@ from typing import Union from typing import overload +import httpx +from tenacity import retry +from tenacity import retry_if_exception_type +from tenacity import stop_after_attempt +from tenacity import wait_random_exponential + from autoblocks._impl import global_state from autoblocks._impl.context_vars import EvaluatorRunContext from autoblocks._impl.context_vars import TestCaseRunContext @@ -52,6 +58,32 @@ DEFAULT_MAX_TEST_CASE_CONCURRENCY = 10 +def create_retry_decorator(retry_count: int) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Create retry decorator following existing tenacity pattern.""" + if retry_count <= 0: + # No retry - return identity decorator + def no_retry_decorator(func: Callable[..., Any]) -> Callable[..., Any]: + return func + + return no_retry_decorator + + # Use same pattern as existing code in api.py + return retry( + stop=stop_after_attempt(retry_count + 1), # +1 because first attempt isn't a retry + wait=wait_random_exponential(multiplier=1, max=30), + retry=retry_if_exception_type( + ( + httpx.TimeoutException, + httpx.ConnectError, + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.PoolTimeout, + ) + ), + reraise=True, + ) + + def tests_and_hashes_overrides_map() -> Optional[dict[str, list[str]]]: """ AUTOBLOCKS_OVERRIDES_TESTS_AND_HASHES is a JSON string that maps test suite IDs to a list of test case hashes. @@ -225,6 +257,7 @@ async def run_test_case( evaluators: Sequence[BaseTestEvaluator], fn: Union[Callable[[TestCaseType], Any], Callable[[TestCaseType], Awaitable[Any]]], before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]], + retry_count: int = 0, ) -> None: reset_token = test_case_run_context_var.set( TestCaseRunContext( @@ -233,14 +266,23 @@ async def run_test_case( test_case_hash=test_case_ctx.hash(), ), ) - try: - output, hook_results, test_case_result_id = await run_test_case_unsafe( + + # Create retry decorator + retry_decorator = create_retry_decorator(retry_count) + + # Wrap the unsafe function with retry logic + @retry_decorator + async def _retry_wrapper() -> Tuple[Any, Any, str]: + return await run_test_case_unsafe( test_id=test_id, run_id=run_id, test_case_ctx=test_case_ctx, fn=fn, before_evaluators_hook=before_evaluators_hook, ) + + try: + output, hook_results, test_case_result_id = await _retry_wrapper() except Exception as err: await send_error( test_id=test_id, @@ -363,6 +405,7 @@ async def run_test_suite_for_grid_combo( grid_search_run_group_id: Optional[str], grid_search_params_combo: Optional[GridSearchParamsCombo], human_review_job: Optional[CreateHumanReviewJob], + retry_count: int = 0, ) -> None: try: # Determine message with priority: unified overrides > legacy env var @@ -402,6 +445,7 @@ async def run_test_suite_for_grid_combo( evaluators=evaluators, fn=fn, before_evaluators_hook=before_evaluators_hook, + retry_count=retry_count, ) for test_case_ctx in yield_test_case_contexts_from_test_cases(test_cases) ], @@ -457,6 +501,7 @@ async def async_run_test_suite( caller_filepath: Optional[str], grid_search_params: Optional[GridSearchParams], human_review_job: Optional[CreateHumanReviewJob], + retry_count: int = 0, ) -> None: # Handle alignment mode align_test_id = AutoblocksEnvVar.ALIGN_TEST_EXTERNAL_ID.get() @@ -532,6 +577,7 @@ async def async_run_test_suite( grid_search_run_group_id=None, grid_search_params_combo=None, human_review_job=human_review_job, + retry_count=retry_count, ) except Exception as err: await send_error( @@ -575,6 +621,7 @@ async def async_run_test_suite( grid_search_run_group_id=grid_search_run_group_id, grid_search_params_combo=grid_params_combo, human_review_job=human_review_job, + retry_count=retry_count, ) for grid_params_combo in yield_grid_search_param_combos(grid_search_params) ], @@ -600,6 +647,7 @@ def run_test_suite( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]] = None, grid_search_params: Optional[GridSearchParams] = None, human_review_job: Optional[CreateHumanReviewJob] = None, + retry_count: int = 0, ) -> None: ... @@ -614,6 +662,7 @@ def run_test_suite( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]] = None, grid_search_params: Optional[GridSearchParams] = None, human_review_job: Optional[CreateHumanReviewJob] = None, + retry_count: int = 0, ) -> None: ... @@ -627,6 +676,7 @@ def run_test_suite( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]] = None, grid_search_params: Optional[GridSearchParams] = None, human_review_job: Optional[CreateHumanReviewJob] = None, + retry_count: int = 0, ) -> None: if not is_cli_running(): log.info(f"Running test suite '{id}'") @@ -649,6 +699,7 @@ def run_test_suite( caller_filepath=caller_filepath, grid_search_params=grid_search_params, human_review_job=human_review_job, + retry_count=retry_count, ), global_state.event_loop(), ).result() diff --git a/autoblocks/_impl/testing/v2/run.py b/autoblocks/_impl/testing/v2/run.py index c14cb1a9..40ef8b32 100644 --- a/autoblocks/_impl/testing/v2/run.py +++ b/autoblocks/_impl/testing/v2/run.py @@ -12,11 +12,16 @@ from typing import Union from typing import overload +import httpx from opentelemetry import trace from opentelemetry.baggage import set_baggage from opentelemetry.context import attach from opentelemetry.context import detach from opentelemetry.context import get_current +from tenacity import retry +from tenacity import retry_if_exception_type +from tenacity import stop_after_attempt +from tenacity import wait_random_exponential from autoblocks._impl import global_state from autoblocks._impl.context_vars import EvaluatorRunContext @@ -56,6 +61,32 @@ DEFAULT_MAX_TEST_CASE_CONCURRENCY = 10 +def create_retry_decorator(retry_count: int) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Create retry decorator following existing tenacity pattern.""" + if retry_count <= 0: + # No retry - return identity decorator + def no_retry_decorator(func: Callable[..., Any]) -> Callable[..., Any]: + return func + + return no_retry_decorator + + # Use same pattern as existing code in api.py + return retry( + stop=stop_after_attempt(retry_count + 1), # +1 because first attempt isn't a retry + wait=wait_random_exponential(multiplier=1, max=30), + retry=retry_if_exception_type( + ( + httpx.TimeoutException, + httpx.ConnectError, + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.PoolTimeout, + ) + ), + reraise=True, + ) + + def tests_and_hashes_overrides_map() -> Optional[dict[str, list[str]]]: """ AUTOBLOCKS_OVERRIDES_TESTS_AND_HASHES is a JSON string that maps test suite IDs to a list of test case hashes. @@ -244,6 +275,7 @@ async def run_test_case( fn: Union[Callable[[TestCaseType], Any], Callable[[TestCaseType], Awaitable[Any]]], before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]], test_case_idx: int, + retry_count: int = 0, ) -> None: reset_token = test_case_run_context_var.set( TestCaseRunContext( @@ -252,9 +284,14 @@ async def run_test_case( test_case_hash=test_case_ctx.hash(), ), ) - try: - log.info(f"Running test case {test_case_idx} for test suite {test_id}") - await run_test_case_unsafe( + + # Create retry decorator + retry_decorator = create_retry_decorator(retry_count) + + # Wrap the unsafe function with retry logic + @retry_decorator + async def _retry_wrapper() -> None: + return await run_test_case_unsafe( test_id=test_id, app_slug=app_slug, test_case_ctx=test_case_ctx, @@ -262,6 +299,10 @@ async def run_test_case( before_evaluators_hook=before_evaluators_hook, evaluators=evaluators, ) + + try: + log.info(f"Running test case {test_case_idx} for test suite {test_id}") + await _retry_wrapper() except Exception as err: log.error(f"Error running test case '{test_case_ctx.hash()}'", exc_info=err) return @@ -355,6 +396,7 @@ async def run_test_suite_for_grid_combo( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]], grid_search_params_combo: Optional[GridSearchParamsCombo], human_review_job: Optional[CreateHumanReviewJob], + retry_count: int = 0, ) -> None: run_id = cuid_generator() start_timestamp = now_rfc3339() @@ -383,6 +425,7 @@ async def run_test_suite_for_grid_combo( fn=fn, before_evaluators_hook=before_evaluators_hook, test_case_idx=test_case_idx, + retry_count=retry_count, ) for test_case_idx, test_case_ctx in enumerate(yield_test_case_contexts_from_test_cases(test_cases)) ], @@ -422,6 +465,7 @@ async def async_run_test_suite( max_test_case_concurrency: int, grid_search_params: Optional[GridSearchParams], human_review_job: Optional[CreateHumanReviewJob], + retry_count: int = 0, ) -> None: # This will be set if the user passed filters to the CLI @@ -475,6 +519,7 @@ async def async_run_test_suite( before_evaluators_hook=before_evaluators_hook, grid_search_params_combo=None, human_review_job=human_review_job, + retry_count=retry_count, ) except Exception as err: log.error(f"Error running test suite '{test_id}'", exc_info=err) @@ -492,6 +537,7 @@ async def async_run_test_suite( before_evaluators_hook=before_evaluators_hook, grid_search_params_combo=grid_params_combo, human_review_job=human_review_job, + retry_count=retry_count, ) for grid_params_combo in yield_grid_search_param_combos(grid_search_params) ], @@ -512,6 +558,7 @@ def run_test_suite( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]] = None, grid_search_params: Optional[GridSearchParams] = None, human_review_job: Optional[CreateHumanReviewJob] = None, + retry_count: int = 0, ) -> None: ... @@ -527,6 +574,7 @@ def run_test_suite( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]] = None, grid_search_params: Optional[GridSearchParams] = None, human_review_job: Optional[CreateHumanReviewJob] = None, + retry_count: int = 0, ) -> None: ... @@ -541,6 +589,7 @@ def run_test_suite( before_evaluators_hook: Optional[Callable[[TestCaseType, Any], Any]] = None, grid_search_params: Optional[GridSearchParams] = None, human_review_job: Optional[CreateHumanReviewJob] = None, + retry_count: int = 0, ) -> None: if not global_state.is_auto_tracer_initialized(): log.error( @@ -561,6 +610,7 @@ def run_test_suite( max_test_case_concurrency=max_test_case_concurrency, grid_search_params=grid_search_params, human_review_job=human_review_job, + retry_count=retry_count, ), global_state.event_loop(), ).result() diff --git a/tests/autoblocks/test_run_test_suite.py b/tests/autoblocks/test_run_test_suite.py index f6f49f50..be819ea4 100644 --- a/tests/autoblocks/test_run_test_suite.py +++ b/tests/autoblocks/test_run_test_suite.py @@ -9,6 +9,7 @@ from typing import Union from unittest import mock +import httpx import pydantic import pytest @@ -3828,3 +3829,278 @@ def test_fn(test_case: MyTestCase) -> Union[Exception, float]: output = results_req["body"]["testCaseOutput"] assert output.startswith("Traceback (most recent call last):") assert "ZeroDivisionError: division by zero" in output + + +def test_retry_on_timeout_errors(httpx_mock): + """Test that retry_count parameter retries test cases on timeout/connection errors.""" + + mock_run_id = str(uuid.uuid4()) + + # Expect start request + expect_cli_post_request( + httpx_mock, + path="/start", + body=dict( + testExternalId="my-test-id", + gridSearchRunGroupId=None, + gridSearchParamsCombo=None, + ), + json=dict(id=mock_run_id), + ) + + # Expect results request (after successful retry) + expect_cli_post_request( + httpx_mock, + path="/results", + body=dict( + testExternalId="my-test-id", + runId=mock_run_id, + testCaseHash="a", + testCaseBody=dict(input="a"), + testCaseOutput="a!", + testCaseDurationMs=ANY_NUMBER, + testCaseRevisionUsage=None, + testCaseHumanReviewInputFields=None, + testCaseHumanReviewOutputFields=None, + datasetItemId=None, + ), + json=dict(id="mock-result-id-1"), + ) + + # Expect end request + expect_cli_post_request( + httpx_mock, + path="/end", + body=dict( + testExternalId="my-test-id", + runId=mock_run_id, + ), + ) + + # Mock the test function to fail twice then succeed + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + if call_count <= 2: + # First two calls fail with timeout + raise httpx.TimeoutException("Request timed out") + return test_case.input + "!" + + run_test_suite( + id="my-test-id", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=3, # Allow up to 3 retries + ) + + # Should have been called 3 times (2 failures + 1 success) + assert call_count == 3 + + # Verify all expected requests were made + requests = httpx_mock.get_requests() + assert len(requests) == 3 # start, results, end + + +def test_retry_exhausted_reports_error(httpx_mock): + """Test that when retries are exhausted, the error is properly reported.""" + + mock_run_id = str(uuid.uuid4()) + + # Expect start request + expect_cli_post_request( + httpx_mock, + path="/start", + body=dict( + testExternalId="my-test-id", + gridSearchRunGroupId=None, + gridSearchParamsCombo=None, + ), + json=dict(id=mock_run_id), + ) + + # Expect error request (after retries exhausted) + expect_cli_post_request( + httpx_mock, + path="/errors", + body=None, # Will verify content below + ) + + # Expect end request + expect_cli_post_request( + httpx_mock, + path="/end", + body=dict( + testExternalId="my-test-id", + runId=mock_run_id, + ), + ) + + # Mock the test function to always fail with timeout + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + raise httpx.TimeoutException("Request timed out") + + run_test_suite( + id="my-test-id", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=2, # Allow up to 2 retries + ) + + # Should have been called 3 times (2 retries + 1 original attempt) + assert call_count == 3 + + # Verify error was reported + requests = httpx_mock.get_requests() + error_req = [r for r in requests if r.url.path == "/errors"][0] + error_req_body = decode_request_body(error_req) + + assert error_req_body["testExternalId"] == "my-test-id" + assert error_req_body["runId"] == mock_run_id + assert error_req_body["testCaseHash"] == "a" + assert error_req_body["evaluatorExternalId"] is None + assert error_req_body["error"]["name"] == "TimeoutException" + assert error_req_body["error"]["message"] == "Request timed out" + + +def test_retry_only_on_specific_exceptions(httpx_mock): + """Test that retry only happens for specific httpx exceptions, not all exceptions.""" + + mock_run_id = str(uuid.uuid4()) + + # Expect start request + expect_cli_post_request( + httpx_mock, + path="/start", + body=dict( + testExternalId="my-test-id", + gridSearchRunGroupId=None, + gridSearchParamsCombo=None, + ), + json=dict(id=mock_run_id), + ) + + # Expect error request (no retry for ValueError) + expect_cli_post_request( + httpx_mock, + path="/errors", + body=None, + ) + + # Expect end request + expect_cli_post_request( + httpx_mock, + path="/end", + body=dict( + testExternalId="my-test-id", + runId=mock_run_id, + ), + ) + + # Mock the test function to fail with ValueError (should not retry) + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + raise ValueError("Not a timeout error") + + run_test_suite( + id="my-test-id", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=3, # Allow up to 3 retries + ) + + # Should have been called only once (no retries for ValueError) + assert call_count == 1 + + # Verify error was reported + requests = httpx_mock.get_requests() + error_req = [r for r in requests if r.url.path == "/errors"][0] + error_req_body = decode_request_body(error_req) + + assert error_req_body["error"]["name"] == "ValueError" + assert error_req_body["error"]["message"] == "Not a timeout error" + + +def test_retry_count_zero_no_retry(httpx_mock): + """Test that retry_count=0 (default) doesn't retry on timeout errors.""" + + mock_run_id = str(uuid.uuid4()) + + # Expect start request + expect_cli_post_request( + httpx_mock, + path="/start", + body=dict( + testExternalId="my-test-id", + gridSearchRunGroupId=None, + gridSearchParamsCombo=None, + ), + json=dict(id=mock_run_id), + ) + + # Expect error request (no retry) + expect_cli_post_request( + httpx_mock, + path="/errors", + body=None, + ) + + # Expect end request + expect_cli_post_request( + httpx_mock, + path="/end", + body=dict( + testExternalId="my-test-id", + runId=mock_run_id, + ), + ) + + # Mock the test function to fail with timeout + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + raise httpx.TimeoutException("Request timed out") + + run_test_suite( + id="my-test-id", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=0, # No retries + ) + + # Should have been called only once (no retries) + assert call_count == 1 + + # Verify error was reported + requests = httpx_mock.get_requests() + error_req = [r for r in requests if r.url.path == "/errors"][0] + error_req_body = decode_request_body(error_req) + + assert error_req_body["error"]["name"] == "TimeoutException" + assert error_req_body["error"]["message"] == "Request timed out" diff --git a/tests/autoblocks/test_run_test_suite_v2_retry.py b/tests/autoblocks/test_run_test_suite_v2_retry.py new file mode 100644 index 00000000..ef77e0b1 --- /dev/null +++ b/tests/autoblocks/test_run_test_suite_v2_retry.py @@ -0,0 +1,257 @@ +import dataclasses +import os +from unittest import mock + +import httpx +import pytest + +from autoblocks._impl.util import AutoblocksEnvVar +from autoblocks.testing.models import BaseTestCase +from autoblocks.testing.models import BaseTestEvaluator +from autoblocks.testing.models import Evaluation +from autoblocks.testing.v2.run import run_test_suite +from autoblocks.tracer import init_auto_tracer + + +@dataclasses.dataclass +class MyTestCase(BaseTestCase): + input: str + + def hash(self) -> str: + return self.input + + +class MyEvaluator(BaseTestEvaluator): + id = "my-evaluator" + + def evaluate_test_case(self, test_case: MyTestCase, output: str) -> Evaluation: + return Evaluation(score=1.0, metadata=dict(output=output)) + + +@pytest.fixture(autouse=True) +def mock_env_vars(): + """Mock environment variables needed for V2 testing.""" + with mock.patch.dict( + os.environ, + { + AutoblocksEnvVar.V2_API_KEY.value: "mock-api-key", + }, + ): + yield + + +@pytest.fixture(autouse=True) +def init_tracer(): + """Initialize auto tracer for V2 tests.""" + init_auto_tracer( + api_key="mock-api-key", + ) + + +def test_v2_retry_on_timeout_errors(): + """Test that retry_count parameter retries test cases on timeout/connection errors in V2.""" + + # Mock the test function to fail twice then succeed + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + if call_count <= 2: + # First two calls fail with timeout + raise httpx.TimeoutException("Request timed out") + return test_case.input + "!" + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=3, # Allow up to 3 retries + ) + + # Should have been called 3 times (2 failures + 1 success) + assert call_count == 3 + + +def test_v2_retry_exhausted_logs_error(): + """Test that when retries are exhausted, the error is properly logged in V2.""" + + # Mock the test function to always fail with timeout + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + raise httpx.TimeoutException("Request timed out") + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=2, # Allow up to 2 retries + ) + + # Should have been called 3 times (2 retries + 1 original attempt) + assert call_count == 3 + + +def test_v2_retry_only_on_specific_exceptions(): + """Test that retry only happens for specific httpx exceptions, not all exceptions in V2.""" + + # Mock the test function to fail with ValueError (should not retry) + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + raise ValueError("Not a timeout error") + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=3, # Allow up to 3 retries + ) + + # Should have been called only once (no retries for ValueError) + assert call_count == 1 + + +def test_v2_retry_count_zero_no_retry(): + """Test that retry_count=0 (default) doesn't retry on timeout errors in V2.""" + + # Mock the test function to fail with timeout + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + raise httpx.TimeoutException("Request timed out") + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=0, # No retries + ) + + # Should have been called only once (no retries) + assert call_count == 1 + + +def test_v2_retry_with_evaluators(): + """Test that retry works correctly with evaluators in V2.""" + + # Mock the test function to fail once then succeed + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + if call_count == 1: + # First call fails with timeout + raise httpx.TimeoutException("Request timed out") + return test_case.input + "!" + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[MyEvaluator()], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=2, # Allow up to 2 retries + ) + + # Should have been called 2 times (1 failure + 1 success) + assert call_count == 2 + + +def test_v2_retry_with_connection_errors(): + """Test that retry works with different types of connection errors in V2.""" + + # Test different types of connection errors + error_types = [ + httpx.ConnectError("Connection failed"), + httpx.ReadTimeout("Read timeout"), + httpx.WriteTimeout("Write timeout"), + httpx.PoolTimeout("Pool timeout"), + ] + + for error_type in error_types: + call_count = 0 + + def test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + if call_count == 1: + # First call fails with the specific error + raise error_type + return test_case.input + "!" + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=test_fn, + max_test_case_concurrency=1, + retry_count=1, # Allow 1 retry + ) + + # Should have been called 2 times (1 failure + 1 success) + assert call_count == 2 + + +def test_v2_retry_with_async_function(): + """Test that retry works with async test functions in V2.""" + + # Mock the async test function to fail once then succeed + call_count = 0 + + async def async_test_fn(test_case: MyTestCase) -> str: + nonlocal call_count + call_count += 1 + if call_count == 1: + # First call fails with timeout + raise httpx.TimeoutException("Request timed out") + return test_case.input + "!" + + run_test_suite( + id="my-test-id", + app_slug="test-app", + test_cases=[ + MyTestCase(input="a"), + ], + evaluators=[], + fn=async_test_fn, + max_test_case_concurrency=1, + retry_count=2, # Allow up to 2 retries + ) + + # Should have been called 2 times (1 failure + 1 success) + assert call_count == 2