From 812704aca1491ac828bf679e754c772c8aa6a48b Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 12 Aug 2026 23:17:21 +0000 Subject: [PATCH] tests(functions): add flaky retries for apply axis 1 tests --- .../tests/system/large/functions/test_managed_function.py | 2 ++ .../tests/system/large/functions/test_remote_function.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/packages/bigframes/tests/system/large/functions/test_managed_function.py b/packages/bigframes/tests/system/large/functions/test_managed_function.py index 888852edd4d1..959b058c8e4a 100644 --- a/packages/bigframes/tests/system/large/functions/test_managed_function.py +++ b/packages/bigframes/tests/system/large/functions/test_managed_function.py @@ -348,6 +348,7 @@ def add_ints(x, y): ) +@pytest.mark.flaky(retries=2, delay=30) def test_managed_function_dataframe_apply_axis_1_array_output( session, dataset_id, function_id ): @@ -560,6 +561,7 @@ def foo(x: int) -> int: )(foo) +@pytest.mark.flaky(retries=2, delay=30) def test_managed_function_df_apply_axis_1( session, dataset_id, scalars_dfs, function_id ): diff --git a/packages/bigframes/tests/system/large/functions/test_remote_function.py b/packages/bigframes/tests/system/large/functions/test_remote_function.py index 69769a1a8461..91886b7919b7 100644 --- a/packages/bigframes/tests/system/large/functions/test_remote_function.py +++ b/packages/bigframes/tests/system/large/functions/test_remote_function.py @@ -2360,6 +2360,7 @@ def foo_named(x: int) -> int: ) +@pytest.mark.flaky(retries=2, delay=120) def test_df_apply_axis_1_multiple_params(session): bf_df = bigframes.dataframe.DataFrame( { @@ -2440,6 +2441,7 @@ def foo(x, y, z): cleanup_function_assets(foo, session.bqclient, session.cloudfunctionsclient) +@pytest.mark.flaky(retries=2, delay=120) def test_df_apply_axis_1_multiple_params_array_output(session): bf_df = bigframes.dataframe.DataFrame( { @@ -2524,6 +2526,7 @@ def foo(x, y, z): cleanup_function_assets(foo, session.bqclient, session.cloudfunctionsclient) +@pytest.mark.flaky(retries=2, delay=120) def test_df_apply_axis_1_single_param_non_series(session): bf_df = bigframes.dataframe.DataFrame( {