diff --git a/automl/snippets/requirements-test.txt b/automl/snippets/requirements-test.txt
index e29e011b390c..81c3a60c109a 100644
--- a/automl/snippets/requirements-test.txt
+++ b/automl/snippets/requirements-test.txt
@@ -1,2 +1,2 @@
-backoff==1.11.1
+backoff==2.0.0
 pytest==7.1.2
diff --git a/automl/tables/predict_test.py b/automl/tables/predict_test.py
index 5530e6eb8dd4..42bfd89cb11b 100644
--- a/automl/tables/predict_test.py
+++ b/automl/tables/predict_test.py
@@ -31,7 +31,9 @@
 
 
 @backoff.on_exception(
-    wait_gen=lambda: iter([50, 150, MAX_TIMEOUT]), exception=Exception, max_tries=3
+    wait_gen=lambda: (wait_time for wait_time in [50, 150, MAX_TIMEOUT]),
+    exception=Exception,
+    max_tries=3,
 )
 def test_predict(capsys):
     inputs = {
diff --git a/automl/tables/requirements-test.txt b/automl/tables/requirements-test.txt
index 6d86ae0febcf..3952c170c352 100644
--- a/automl/tables/requirements-test.txt
+++ b/automl/tables/requirements-test.txt
@@ -1,2 +1,2 @@
 pytest==7.1.2
-backoff==1.11.1
\ No newline at end of file
+backoff==2.0.0
\ No newline at end of file