From 4da767d17c85e10023478c2dba25e90bcb70f3f1 Mon Sep 17 00:00:00 2001
From: Shaharyar Shamshi <shaharyarshamshi@gmail.com>
Date: Sun, 3 Mar 2024 05:00:44 +0530
Subject: [PATCH] test case improved

---
 .../DateRangeCalendar.test.tsx                | 30 +++++++++++-----
 .../tests/timezone.DateCalendar.test.tsx      | 36 +++++++++++++------
 2 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/packages/x-date-pickers-pro/src/DateRangeCalendar/DateRangeCalendar.test.tsx b/packages/x-date-pickers-pro/src/DateRangeCalendar/DateRangeCalendar.test.tsx
index ed000144da530..1a4215ffa5271 100644
--- a/packages/x-date-pickers-pro/src/DateRangeCalendar/DateRangeCalendar.test.tsx
+++ b/packages/x-date-pickers-pro/src/DateRangeCalendar/DateRangeCalendar.test.tsx
@@ -573,18 +573,32 @@ describe('<DateRangeCalendar />', () => {
         // Render the component with initial timezone prop
         const { rerender } = render(<DateRangeCalendar timezone="UTC" />);
 
-        const renderButtons = screen
-          .getAllByRole('gridcell')
-          .filter((button) => button.tagName.toLowerCase() === 'button');
+        // Create a map of buttons with their indices for the initial render
+        const renderButtonsMap = {};
+        Object.keys(screen.getAllByRole('gridcell')).forEach((key, index) => {
+          const element = screen.getByRole('gridcell', { name: key });
+          if (element.tagName.toLowerCase() === 'button') {
+            renderButtonsMap[index] = element;
+          }
+        });
 
+        // Rerender the component with a different timezone prop
         rerender(<DateRangeCalendar timezone={timezone} />);
 
-        const reRenderButtons = screen
-          .getAllByRole('gridcell')
-          .filter((button) => button.tagName.toLowerCase() === 'button');
+        // Create a map of buttons with their indices for the rerender
+        const reRenderButtonsMap = {};
+        Object.keys(screen.getAllByRole('gridcell')).forEach((key, index) => {
+          const element = screen.getByRole('gridcell', { name: key });
+          if (element.tagName.toLowerCase() === 'button') {
+            reRenderButtonsMap[index] = element;
+          }
+        });
 
-        // check the number of the button rendered equal
-        expect(reRenderButtons.length).equals(renderButtons.length);
+        // Ensure the number of buttons and their innerText are consistent
+        expect(Object.keys(reRenderButtonsMap).length).equals(Object.keys(renderButtonsMap).length);
+        Object.keys(renderButtonsMap).forEach((index) => {
+          expect(reRenderButtonsMap[index].innerText).equals(renderButtonsMap[index].innerText);
+        });
       });
 
       it('should select the range from the next month', () => {
diff --git a/packages/x-date-pickers/src/DateCalendar/tests/timezone.DateCalendar.test.tsx b/packages/x-date-pickers/src/DateCalendar/tests/timezone.DateCalendar.test.tsx
index 0d08c8a156e34..370bee39ce0d3 100644
--- a/packages/x-date-pickers/src/DateCalendar/tests/timezone.DateCalendar.test.tsx
+++ b/packages/x-date-pickers/src/DateCalendar/tests/timezone.DateCalendar.test.tsx
@@ -65,18 +65,34 @@ describe('<DateCalendar /> - Timezone', () => {
           // Render the component with initial timezone prop
           const { rerender } = render(<DateCalendar timezone="UTC" />);
 
-          const renderButtons = screen
-            .getAllByRole('gridcell')
-            .filter((button) => button.tagName.toLowerCase() === 'button');
-
+          // Create a map of buttons with their indices for the initial render
+          const renderButtonsMap = {};
+          Object.keys(screen.getAllByRole('gridcell')).forEach((key, index) => {
+            const element = screen.getByRole('gridcell', { name: key });
+            if (element.tagName.toLowerCase() === 'button') {
+              renderButtonsMap[index] = element;
+            }
+          });
+
+          // Rerender the component with a different timezone prop
           rerender(<DateCalendar timezone={timezone} />);
 
-          const reRenderButtons = screen
-            .getAllByRole('gridcell')
-            .filter((button) => button.tagName.toLowerCase() === 'button');
-
-          // check the number of the button rendered equal
-          expect(reRenderButtons.length).equals(renderButtons.length);
+          // Create a map of buttons with their indices for the rerender
+          const reRenderButtonsMap = {};
+          Object.keys(screen.getAllByRole('gridcell')).forEach((key, index) => {
+            const element = screen.getByRole('gridcell', { name: key });
+            if (element.tagName.toLowerCase() === 'button') {
+              reRenderButtonsMap[index] = element;
+            }
+          });
+
+          // Ensure the number of buttons and their innerText are consistent
+          expect(Object.keys(reRenderButtonsMap).length).equals(
+            Object.keys(renderButtonsMap).length,
+          );
+          Object.keys(renderButtonsMap).forEach((index) => {
+            expect(reRenderButtonsMap[index].innerText).equals(renderButtonsMap[index].innerText);
+          });
         });
 
         it('on timezone change the difference between the selected date change should be at max 1', () => {