diff --git a/test/test_lrg_2d.py b/test/test_lrg_2d.py
index f6b1d369..8a48b9f2 100644
--- a/test/test_lrg_2d.py
+++ b/test/test_lrg_2d.py
@@ -69,7 +69,7 @@ def fake_gen_zarr(file_path, input_fname):
         x_no_api=True,
         return_state=True,
     )
-    assert state.is_failed()
+    assert state.is_completed(), "Flow failed"
 
     response = {}
     with open(mock_callback_data) as fd:
@@ -105,9 +105,7 @@ def fake_gen_zarr(file_path, input_fname):
         x_no_api=True,
         return_state=True,
     )
-    assert (
-        state.is_failed()
-    )  # Except for BRT, everything else should be failed (although partly failed)
+    assert state.is_completed(), "Flow failed"
 
     response = {}
     with open(mock_callback_data) as fd: