diff --git a/tests/container-images/gerrit-init/test_container_integration_gerrit_init.py b/tests/container-images/gerrit-init/test_container_integration_gerrit_init.py
index 287dac7..9de4043 100644
--- a/tests/container-images/gerrit-init/test_container_integration_gerrit_init.py
+++ b/tests/container-images/gerrit-init/test_container_integration_gerrit_init.py
@@ -91,17 +91,7 @@
 
     @pytest.mark.timeout(60)
     def test_gerrit_init_exits_after_init(self, container_run_default):
-        def wait_for_container_exit():
-            try:
-                container_run_default.reload()
-                return False
-            except NotFound:
-                return True
-
-        while not wait_for_container_exit():
-            continue
-
-        assert container_run_default.attrs["State"]["ExitCode"] == 0
+        assert container_run_default.wait()["StatusCode"] == 0
 
 
 @pytest.fixture(