From 0a657ee3bed0468478fa6d7823bdb509bd4aa172 Mon Sep 17 00:00:00 2001
From: haseebmalik18 <haseebulhaq55@gmail.com>
Date: Sun, 5 Apr 2026 15:12:03 -0400
Subject: [PATCH] Improve Playwright test patterns in dag-calendar-tab #63409

---
 .../ui/tests/e2e/pages/DagCalendarTab.ts      | 26 +++++--------------
 .../tests/e2e/specs/dag-calendar-tab.spec.ts  | 10 +++----
 2 files changed, 10 insertions(+), 26 deletions(-)

diff --git a/airflow-core/src/airflow/ui/tests/e2e/pages/DagCalendarTab.ts b/airflow-core/src/airflow/ui/tests/e2e/pages/DagCalendarTab.ts
index 7aee20db7c083..d3dbc1cf6ddfb 100644
--- a/airflow-core/src/airflow/ui/tests/e2e/pages/DagCalendarTab.ts
+++ b/airflow-core/src/airflow/ui/tests/e2e/pages/DagCalendarTab.ts
@@ -58,10 +58,6 @@ export class DagCalendarTab extends BasePage {
     return colors;
   }
 
-  public async getActiveCellCount(): Promise<number> {
-    return this.activeCells.count();
-  }
-
   public async getManualRunStates(): Promise<Array<string>> {
     const count = await this.activeCells.count();
     const states: Array<string> = [];
@@ -88,7 +84,7 @@ export class DagCalendarTab extends BasePage {
   public async navigateToCalendar(dagId: string) {
     await expect(async () => {
       await this.safeGoto(`/dags/${dagId}/calendar`);
-      await this.page.getByTestId("dag-calendar-root").waitFor({ state: "visible", timeout: 5000 });
+      await expect(this.page.getByTestId("dag-calendar-root")).toBeVisible({ timeout: 5000 });
     }).toPass({ intervals: [2000], timeout: 60_000 });
     await this.waitForCalendarReady();
   }
@@ -100,7 +96,7 @@ export class DagCalendarTab extends BasePage {
   public async switchToHourly() {
     await this.hourlyToggle.click();
 
-    await this.page.getByTestId("calendar-hourly-view").waitFor({ state: "visible", timeout: 30_000 });
+    await expect(this.page.getByTestId("calendar-hourly-view")).toBeVisible({ timeout: 30_000 });
   }
 
   public async switchToTotalView() {
@@ -112,22 +108,14 @@ export class DagCalendarTab extends BasePage {
   }
 
   private async waitForCalendarReady(): Promise<void> {
-    await this.page.getByTestId("dag-calendar-root").waitFor({ state: "visible", timeout: 120_000 });
-
-    await this.page.getByTestId("calendar-current-period").waitFor({ state: "visible", timeout: 120_000 });
+    await expect(this.page.getByTestId("dag-calendar-root")).toBeVisible({ timeout: 120_000 });
 
-    const overlay = this.page.getByTestId("calendar-loading-overlay");
-
-    if (await overlay.isVisible().catch(() => false)) {
-      await overlay.waitFor({ state: "hidden", timeout: 120_000 });
-    }
+    await expect(this.page.getByTestId("calendar-current-period")).toBeVisible({ timeout: 120_000 });
 
-    await this.page.getByTestId("calendar-grid").waitFor({ state: "visible", timeout: 120_000 });
+    await expect(this.page.getByTestId("calendar-loading-overlay")).toBeHidden({ timeout: 120_000 });
 
-    await this.page.waitForFunction(() => {
-      const cells = document.querySelectorAll('[data-testid="calendar-cell"]');
+    await expect(this.page.getByTestId("calendar-grid")).toBeVisible({ timeout: 120_000 });
 
-      return cells.length > 0;
-    });
+    await expect(this.calendarCells.first()).toBeVisible({ timeout: 120_000 });
   }
 }
diff --git a/airflow-core/src/airflow/ui/tests/e2e/specs/dag-calendar-tab.spec.ts b/airflow-core/src/airflow/ui/tests/e2e/specs/dag-calendar-tab.spec.ts
index 65d77ac80d0a5..c429045741ca6 100644
--- a/airflow-core/src/airflow/ui/tests/e2e/specs/dag-calendar-tab.spec.ts
+++ b/airflow-core/src/airflow/ui/tests/e2e/specs/dag-calendar-tab.spec.ts
@@ -104,9 +104,7 @@ test.describe("DAG Calendar Tab", () => {
   test("verify active cells appear for DAG runs", async () => {
     await calendar.switchToHourly();
 
-    const count = await calendar.getActiveCellCount();
-
-    expect(count).toBeGreaterThan(0);
+    await expect(calendar.activeCells.first()).toBeVisible();
   });
 
   test("verify manual runs are detected", async () => {
@@ -150,13 +148,11 @@ test.describe("DAG Calendar Tab", () => {
   test("failed view reduces active cells", async () => {
     await calendar.switchToHourly();
 
-    const totalCount = await calendar.getActiveCellCount();
+    const totalCount = await calendar.activeCells.count();
 
     await calendar.switchToFailedView();
 
-    const failedCount = await calendar.getActiveCellCount();
-
-    expect(failedCount).toBeLessThan(totalCount);
+    await expect.poll(async () => calendar.activeCells.count(), { timeout: 10_000 }).toBeLessThan(totalCount);
   });
 
   test.fixme("color scale changes between total and failed view", async () => {