getCellOutput is null in widget ui-test

Widgets
1

In my ui-test, the cell output is null,

    await page.notebook.runCellByCell({
      onAfterCellRun: async (cellIndex: number) => {
        const cell = await page.notebook.getCellOutput(cellIndex);
        if (cell) {
          captures.push(await cell.screenshot());
        }
      },
    });

I then try to wait for the output for 60 seconds, which is far enough for the cell output to be ready.

await page.notebook.runCellByCell({
      onAfterCellRun: async (cellIndex: number) => {
        let cell = await page.notebook.getCellOutput(cellIndex);
        const startTime = Date.now();
        const timeout = 60000; // Timeout in milliseconds, adjust as necessary

        // Polling for cell output to be not null
        while (!cell && Date.now() - startTime < timeout) {
          await page.waitForTimeout(3000); // Wait for 3 second before retrying  
          console.log("waiting for cell output to be not null: ", cellIndex);
          cell = await page.notebook.getCellOutput(cellIndex);
        }

        if (cell) {
          captures.push(await cell.screenshot());
        } else {
          console.log("Cell output is not available for cell:", cellIndex);
        }
      },
    });

but sometimes, the cell output is still null, and the test passes randomly.
Any idea how to fix it?