test: use expect.poll instead of untilUpdated (#511)

sapphi-red · web-flow · commit 18e89ce28384 · 2025-06-17T09:48:49.000+09:00
diff --git a/playground/class-components/__tests__/class-components.spec.ts b/playground/class-components/__tests__/class-components.spec.ts
@@ -1,11 +1,5 @@
 import { expect, test } from 'vitest'
-import {
-  editFile,
-  isServe,
-  page,
-  untilBrowserLogAfter,
-  untilUpdated,
-} from '~utils'
+import { editFile, isServe, page, untilBrowserLogAfter } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('span')).toMatch('Hello World')
@@ -18,13 +12,17 @@ if (isServe) {
       () => page.textContent('span'),
       '[vite] hot updated: /src/App.tsx',
     )
-    await untilUpdated(() => page.textContent('span'), 'Hello class components')
+    await expect
+      .poll(() => page.textContent('span'))
+      .toMatch('Hello class components')
 
     editFile('src/utils.tsx', (code) => code.replace('Hello', 'Hi'))
     await untilBrowserLogAfter(
       () => page.textContent('span'),
       '[vite] hot updated: /src/App.tsx',
     )
-    await untilUpdated(() => page.textContent('span'), 'Hi class components')
+    await expect
+      .poll(() => page.textContent('span'))
+      .toMatch('Hi class components')
   })
 }
diff --git a/playground/compiler-react-18/__tests__/compiler-react-18.spec.ts b/playground/compiler-react-18/__tests__/compiler-react-18.spec.ts
@@ -1,5 +1,5 @@
 import { expect, test } from 'vitest'
-import { editFile, isServe, page, untilUpdated } from '~utils'
+import { editFile, isServe, page } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('button')).toMatch('count is 0')
@@ -11,5 +11,5 @@ test.runIf(isServe)('should hmr', async () => {
   editFile('src/App.tsx', (code) =>
     code.replace('count is {count}', 'count is {count}!'),
   )
-  await untilUpdated(() => page.textContent('button'), 'count is 1!')
+  await expect.poll(() => page.textContent('button')).toMatch('count is 1!')
 })
diff --git a/playground/compiler/__tests__/compiler.spec.ts b/playground/compiler/__tests__/compiler.spec.ts
@@ -1,5 +1,5 @@
 import { expect, test } from 'vitest'
-import { editFile, isServe, page, untilUpdated } from '~utils'
+import { editFile, isServe, page } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('button')).toMatch('count is 0')
@@ -11,5 +11,5 @@ test.runIf(isServe)('should hmr', async () => {
   editFile('src/App.tsx', (code) =>
     code.replace('count is {count}', 'count is {count}!'),
   )
-  await untilUpdated(() => page.textContent('button'), 'count is 1!')
+  await expect.poll(() => page.textContent('button')).toMatch('count is 1!')
 })
diff --git a/playground/hook-with-jsx/__tests__/hook-with-jsx.spec.ts b/playground/hook-with-jsx/__tests__/hook-with-jsx.spec.ts
@@ -1,11 +1,5 @@
 import { expect, test } from 'vitest'
-import {
-  editFile,
-  isServe,
-  page,
-  untilBrowserLogAfter,
-  untilUpdated,
-} from '~utils'
+import { editFile, isServe, page, untilBrowserLogAfter } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('button')).toMatch('count is 0')
@@ -22,6 +16,6 @@ if (isServe) {
       () => page.textContent('button'),
       '[vite] hot updated: /src/App.tsx',
     )
-    await untilUpdated(() => page.textContent('button'), 'count is 1!')
+    await expect.poll(() => page.textContent('button')).toMatch('count is 1!')
   })
 }
diff --git a/playground/mdx/__tests__/mdx.spec.ts b/playground/mdx/__tests__/mdx.spec.ts
@@ -1,11 +1,5 @@
 import { expect, test } from 'vitest'
-import {
-  editFile,
-  isServe,
-  page,
-  untilBrowserLogAfter,
-  untilUpdated,
-} from '~utils'
+import { editFile, isServe, page, untilBrowserLogAfter } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('h1')).toMatch('Vite + MDX')
@@ -24,7 +18,7 @@ if (isServe) {
       () => page.textContent('h1'),
       '[vite] hot updated: /src/demo.mdx',
     )
-    await untilUpdated(() => page.textContent('h1'), 'Updated')
+    await expect.poll(() => page.textContent('h1')).toMatch('Updated')
   })
 
   test('should hmr with .md extension', async () => {
@@ -35,9 +29,8 @@ if (isServe) {
         ),
       '[vite] hot updated: /src/demo2.md',
     )
-    await untilUpdated(
-      () => page.getByText('.md extension hmr works.').textContent(),
-      '.md extension hmr works. This is bold text.',
-    )
+    await expect
+      .poll(() => page.getByText('.md extension hmr works.').textContent())
+      .toMatch('.md extension hmr works. This is bold text.')
   })
 }
diff --git a/playground/react-classic/__tests__/react.spec.ts b/playground/react-classic/__tests__/react.spec.ts
@@ -1,5 +1,5 @@
 import { expect, test } from 'vitest'
-import { editFile, isServe, page, untilUpdated, viteTestUrl } from '~utils'
+import { editFile, isServe, page, viteTestUrl } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('h1')).toMatch('Hello Vite + React')
@@ -13,7 +13,7 @@ test('should update', async () => {
 
 test.runIf(isServe)('should hmr', async () => {
   editFile('App.jsx', (code) => code.replace('Vite + React', 'Updated'))
-  await untilUpdated(() => page.textContent('h1'), 'Hello Updated')
+  await expect.poll(() => page.textContent('h1')).toMatch('Hello Updated')
   // preserve state
   expect(await page.textContent('button')).toMatch('count is: 1')
 })
diff --git a/playground/react-emotion/__tests__/react.spec.ts b/playground/react-emotion/__tests__/react.spec.ts
@@ -1,5 +1,5 @@
 import { expect, test } from 'vitest'
-import { editFile, getColor, isServe, page, untilUpdated } from '~utils'
+import { editFile, getColor, isServe, page } from '~utils'
 
 test('should render', async () => {
   expect(await page.textContent('h1')).toMatch(
@@ -17,13 +17,13 @@ test.runIf(isServe)('should hmr', async () => {
   editFile('App.jsx', (code) =>
     code.replace('Vite + React + @emotion/react', 'Updated'),
   )
-  await untilUpdated(() => page.textContent('h1'), 'Hello Updated')
+  await expect.poll(() => page.textContent('h1')).toMatch('Hello Updated')
 
   editFile('Counter.jsx', (code) =>
     code.replace('color: #646cff;', 'color: #d26ac2;'),
   )
 
-  await untilUpdated(() => getColor('code'), '#d26ac2')
+  await expect.poll(() => getColor('code')).toMatch('#d26ac2')
 
   // preserve state
   expect(await page.textContent('button')).toMatch('count is: 1')
@@ -47,7 +47,7 @@ test('should update button style', async () => {
       code.replace('border: 2px solid #000', 'border: 4px solid red'),
     )
 
-    await untilUpdated(getButtonBorderStyle, '4px solid rgb(255, 0, 0)')
+    await expect.poll(getButtonBorderStyle).toMatch('4px solid rgb(255, 0, 0)')
 
     // preserve state
     expect(await page.textContent('button')).toMatch('count is: 1')
diff --git a/playground/react/__tests__/react.spec.ts b/playground/react/__tests__/react.spec.ts
@@ -6,7 +6,6 @@ import {
   isServe,
   page,
   untilBrowserLogAfter,
-  untilUpdated,
   viteTestUrl,
 } from '~utils'
 
@@ -24,25 +23,26 @@ test.runIf(isServe)('should hmr', async () => {
   editFile('App.jsx', (code) =>
     code.replace('Vite + React', 'Vite + React Updated'),
   )
-  await untilUpdated(() => page.textContent('h1'), 'Hello Vite + React Updated')
+  await expect
+    .poll(() => page.textContent('h1'))
+    .toMatch('Hello Vite + React Updated')
   // preserve state
   expect(await page.textContent('#state-button')).toMatch('count is: 1')
 
   editFile('App.jsx', (code) =>
     code.replace('Vite + React Updated', 'Vite + React'),
   )
-  await untilUpdated(() => page.textContent('h1'), 'Hello Vite + React')
+  await expect.poll(() => page.textContent('h1')).toMatch('Hello Vite + React')
 })
 
 test.runIf(isServe)('should not invalidate when code is invalid', async () => {
   editFile('App.jsx', (code) =>
     code.replace('<div className="App">', '<div className="App"}>'),
   )
 
-  await untilUpdated(
-    () => page.textContent('vite-error-overlay .message-body'),
-    'Unexpected token',
-  )
+  await expect
+    .poll(() => page.textContent('vite-error-overlay .message-body'))
+    .toMatch('Unexpected token')
   // if import.meta.invalidate happened, the old page won't be shown because the page is reloaded
   expect(await page.textContent('h1')).toMatch('Hello Vite + React')
 
@@ -90,7 +90,7 @@ if (!isBuild) {
         'Parent rendered',
       ],
     )
-    await untilUpdated(() => page.textContent('#parent'), 'Updated')
+    await expect.poll(() => page.textContent('#parent')).toMatch('Updated')
   })
 
   // #3301
@@ -122,10 +122,9 @@ if (!isBuild) {
         'Parent rendered',
       ],
     )
-    await untilUpdated(
-      () => page.textContent('#context-provider'),
-      'context provider updated',
-    )
+    await expect
+      .poll(() => page.textContent('#context-provider'))
+      .toMatch('context provider updated')
   })
 
   test('should hmr files with "react/jsx-runtime"', async () => {
@@ -143,10 +142,9 @@ if (!isBuild) {
         ),
       ['[vite] hot updated: /hmr/jsx-import-runtime.js'],
     )
-    await untilUpdated(
-      () => page.textContent('#jsx-import-runtime'),
-      'JSX import runtime updated',
-    )
+    await expect
+      .poll(() => page.textContent('#jsx-import-runtime'))
+      .toMatch('JSX import runtime updated')
 
     expect(await page.textContent('#state-button')).toMatch('count is: 1')
   })
diff --git a/playground/ssr-react/__tests__/ssr-react.spec.ts b/playground/ssr-react/__tests__/ssr-react.spec.ts
@@ -6,7 +6,6 @@ import {
   isBuild,
   page,
   untilBrowserLogAfter,
-  untilUpdated,
   viteTestUrl as url,
 } from '~utils'
 
@@ -51,11 +50,11 @@ test('/', async () => {
 test.skipIf(isBuild)('hmr', async () => {
   await untilBrowserLogAfter(() => page.goto(url), 'hydrated')
 
-  await untilUpdated(() => page.textContent('h1'), 'Home')
+  await expect.poll(() => page.textContent('h1')).toMatch('Home')
   editFile('src/pages/Home.jsx', (code) =>
     code.replace('<h1>Home', '<h1>changed'),
   )
-  await untilUpdated(() => page.textContent('h1'), 'changed')
+  await expect.poll(() => page.textContent('h1')).toMatch('changed')
 
   // verify the change also affects next SSR
   const res = await page.reload()
@@ -65,14 +64,14 @@ test.skipIf(isBuild)('hmr', async () => {
 test('client navigation', async () => {
   await untilBrowserLogAfter(() => page.goto(url), 'hydrated')
 
-  await untilUpdated(() => page.textContent('a[href="/about"]'), 'About')
+  await expect.poll(() => page.textContent('a[href="/about"]')).toMatch('About')
   await page.click('a[href="/about"]')
-  await untilUpdated(() => page.textContent('h1'), 'About')
+  await expect.poll(() => page.textContent('h1')).toMatch('About')
 
   if (!isBuild) {
     editFile('src/pages/About.jsx', (code) =>
       code.replace('<h1>About', '<h1>changed'),
     )
-    await untilUpdated(() => page.textContent('h1'), 'changed')
+    await expect.poll(() => page.textContent('h1')).toMatch('changed')
   }
 })
diff --git a/playground/test-utils.ts b/playground/test-utils.ts
@@ -37,8 +37,6 @@ function rgbToHex(rgb: string): string {
   }
 }
 
-const timeout = (n: number) => new Promise((r) => setTimeout(r, n))
-
 async function toEl(el: string | ElementHandle): Promise<ElementHandle> {
   if (typeof el === 'string') {
     return await page.$(el)
@@ -84,25 +82,6 @@ export function removeFile(filename: string): void {
   fs.unlinkSync(path.resolve(testDir, filename))
 }
 
-/**
- * Poll a getter until the value it returns includes the expected value.
- */
-export async function untilUpdated(
-  poll: () => string | Promise<string>,
-  expected: string,
-): Promise<void> {
-  const maxTries = process.env.CI ? 100 : 50
-  for (let tries = 0; tries < maxTries; tries++) {
-    const actual = (await poll()) ?? ''
-    if (actual.indexOf(expected) > -1 || tries === maxTries - 1) {
-      expect(actual).toMatch(expected)
-      break
-    } else {
-      await timeout(50)
-    }
-  }
-}
-
 type UntilBrowserLogAfterCallback = (logs: string[]) => PromiseLike<void> | void
 
 export async function untilBrowserLogAfter(
diff --git a/playground/vitest.config.e2e.ts b/playground/vitest.config.e2e.ts
@@ -17,6 +17,11 @@ export default defineConfig({
     testTimeout: timeout,
     hookTimeout: timeout,
     reporters: 'dot',
+    expect: {
+      poll: {
+        timeout: 50 * (process.env.CI ? 200 : 50),
+      },
+    },
   },
   esbuild: {
     target: 'node14',
