diff --git a/tests/src/ExistingSiteJavascript/CollaboraIntegrationTest.php b/tests/src/ExistingSiteJavascript/CollaboraIntegrationTest.php
index f3ba3f8d..a6bc4b6a 100644
--- a/tests/src/ExistingSiteJavascript/CollaboraIntegrationTest.php
+++ b/tests/src/ExistingSiteJavascript/CollaboraIntegrationTest.php
@@ -34,51 +34,23 @@ public function testCollaboraPreview(): void {
         $this->drupalGet('/cool/view/' . $media->id());
         $this->getSession()->switchToIFrame('collabora-online-viewer');
 
-        $this->waitUntilNoMessage(function (): string|null {
-            $canvas = $this->getCurrentPage()->find('css', 'canvas#document-canvas');
-            if (!$canvas) {
-                return 'Canvas element not found.';
-            }
-            return NULL;
-        });
+        $assert_session = $this->assertSession();
+        $canvas = $assert_session->waitForElement('css', 'canvas#document-canvas');
+        $this->assertNotNull($canvas, 'The canvas element was not found after 10 seconds.');
 
-        // Make sure the correct document was opened.
-        // Check the document name at the top of the editor.
-        // Wait until the document name element appears and has non-empty text.
-        $this->waitUntilNoMessage(function(): string|null {
-            // Get a fresh element reference in each iteration, to avoid a
-            // StaleElementReference exception.
-            $element = $this->getCurrentPage()->find('css', 'input#document-name-input');
-            if (!$element) {
-                return 'Document name element not found.';
-            }
-            $text = $element->getValue();
-            if (!$text) {
-                return 'Document name is empty: ' . var_export($text, TRUE);
-            }
-            $this->assertSame('shopping-list.txt', $text);
-            return NULL;
+        $document_field = $assert_session->waitForElement('css', 'input#document-name-input');
+        $this->assertNotNull($document_field, 'The document name input was not found after 10 seconds.');
+        $this->getCurrentPage()->waitFor(10, function() use ($document_field) {
+          return $document_field->getValue() === 'shopping-list.txt';
         });
+        $this->assertEquals('shopping-list.txt', $document_field->getValue(), 'The document name input did not contain the correct value after 10 seconds.');
 
-        // The document text is in a canvas element, so instead we compare the
-        // word count and character count.
-        // Wait until the word counter element appears and has a value.
-        $this->waitUntilNoMessage(
-            fn (): string|null => $this->retryIfStaleElement(
-                function (): string|null {
-                    $element = $this->getCurrentPage()->find('css', 'div#StateWordCount');
-                    if (!$element) {
-                        return 'Word count element not found.';
-                    }
-                    $count_string = $element->getText();
-                    if (!$count_string) {
-                        return 'Word count string is empty.';
-                    }
-                    $this->assertSame('2 words, 18 characters', $count_string);
-                    return NULL;
-                },
-            ),
-        );
+        $word_count_element = $assert_session->waitForElement('css', 'div#StateWordCount');
+        $this->assertNotNull($word_count_element, 'The word count element was not found after 10 seconds.');
+        $this->getCurrentPage()->waitFor(10, function() use ($word_count_element) {
+          return $word_count_element->getText() === '2 words, 18 characters';
+        });
+        $this->assertEquals('2 words, 18 characters', $word_count_element->getText(), 'The word count element did not contain the correct text after 10 seconds.');
     }
 
     /**
@@ -117,54 +89,4 @@ protected function createDocumentMedia(string $media_name, string $file_basename
         return $media;
     }
 
-    /**
-     * Calls a callback until it does not throw a StaleElementReference.
-     *
-     * @template T
-     *
-     * @param callable(): T $callback
-     *   Callback to call and retry.
-     * @param int $max_calls
-     *   Maximum number of times the callback should be called.
-     *
-     * @return T
-     *   Return value from the last call to the callback.
-     */
-    protected function retryIfStaleElement(callable $callback, int $max_calls = 10): mixed {
-        for ($i = 1; $i < $max_calls; ++$i) {
-            try {
-                return $callback();
-            }
-            catch (StaleElementReference) {
-                // Next iteration.
-            }
-        }
-        // Call it one more time, without try/catch.
-        return $callback();
-    }
-
-    /**
-     * Waits until a callback returns NULL.
-     *
-     * @param callable(): (string|null) $callback
-     *   Callback that is called in each iteration.
-     *   It should return NULL to wait no longer, or a string message to wait
-     *   and start another iteration.
-     * @param int|float $max_seconds
-     *   Maximum seconds of wait time.
-     */
-    protected function waitUntilNoMessage(callable $callback, int|float $max_seconds = 10): void {
-        $start = microtime(TRUE);
-        $end = $start + $max_seconds;
-        do {
-            $message = $callback();
-            if ($message === NULL) {
-                return;
-            }
-            usleep(10000);
-        } while (microtime(TRUE) < $end);
-
-        $this->fail('Timeout: ' . $message);
-    }
-
 }