Alerting/Chore: Fix TimeRangeInput not working across multiple months (#93622)

packages/grafana-ui/src/components/DateTimePickers/TimeRangeInput.test.tsx

@@ -0,0 +1,55 @@
+import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+
+import { dateTime } from '@grafana/data';
+
+import { TimeRangeInput } from './TimeRangeInput';
+
+describe('TimeRangeInput', () => {
+  // TODO: This test is evergreen - the check that we haven't accidentally closed
+  // the picker still passes without the appropriate fix
+  // Seems to be related to jest-dom and how it handles clicking outside the node etc.
+  it('handles selecting dates over multiple months', async () => {
+    const user = userEvent.setup();
+    const from = dateTime('2024-01-01T00:00:00Z');
+    const to = dateTime('2024-01-01T00:00:00Z');
+    const onChange = jest.fn();
+
+    render(
+      <TimeRangeInput
+        timeZone="utc"
+        onChange={(payload) => {
+          const { from, to } = payload;
+          onChange({ from: from.toString(), to: to.toString() });
+        }}
+        value={{
+          from,
+          to,
+          raw: {
+            from,
+            to,
+          },
+        }}
+      />
+    );
+
+    // TimeRangeInput renders as a button that looks like an input -
+    // the only one we can see at the start is the button to open the picker
+    await user.click(screen.getByRole('button'));
+
+    const [firstOpenCalendarButton] = await screen.findAllByRole('button', { name: /open calendar/i });
+    await user.click(firstOpenCalendarButton);
+
+    // Select two dates that are on different "screens" of the calendar picker - this is where the bug would occur
+    await user.click(await screen.findByLabelText(/january 1, 2024/i));
+    await user.click(await screen.findByLabelText(/next month/i));
+    await user.click(await screen.findByLabelText(/february 28, 2024/i));
+
+    await user.click(await screen.findByText(/apply time range/i));
+
+    expect(onChange).toHaveBeenCalledWith({
+      from: 'Mon Jan 01 2024 00:00:00 GMT+0000',
+      to: 'Wed Feb 28 2024 23:59:59 GMT+0000',
+    });
+  });
+});

packages/grafana-ui/src/components/DateTimePickers/TimeRangeInput.tsx

@@ -1,11 +1,13 @@
 import { css, cx } from '@emotion/css';
-import { FormEvent, MouseEvent, useState } from 'react';
+import { useDialog } from '@react-aria/dialog';
+import { FocusScope } from '@react-aria/focus';
+import { useOverlay } from '@react-aria/overlays';
+import { createRef, FormEvent, MouseEvent, useState } from 'react';
 
 import { dateTime, getDefaultTimeRange, GrafanaTheme2, TimeRange, TimeZone } from '@grafana/data';
 import { selectors } from '@grafana/e2e-selectors';
 
 import { useStyles2 } from '../../themes/ThemeContext';
-import { ClickOutsideWrapper } from '../ClickOutsideWrapper/ClickOutsideWrapper';
 import { Icon } from '../Icon/Icon';
 import { getInputStyles } from '../Input/Input';
 
@@ -77,6 +79,22 @@ export const TimeRangeInput = ({
     onChange({ from, to, raw: { from, to } });
   };
 
+  const overlayRef = createRef<HTMLElement>();
+  const buttonRef = createRef<HTMLButtonElement>();
+
+  const { dialogProps } = useDialog({}, overlayRef);
+
+  const { overlayProps } = useOverlay(
+    {
+      onClose,
+      isDismissable: true,
+      isOpen,
+      shouldCloseOnInteractOutside: (element) => {
+        return !buttonRef.current?.contains(element);
+      },
+    },
+    overlayRef
+  );
   return (
     <div className={styles.container}>
       <button
@@ -84,6 +102,7 @@ export const TimeRangeInput = ({
         className={styles.pickerInput}
         data-testid={selectors.components.TimePicker.openButton}
         onClick={onOpen}
+        ref={buttonRef}
       >
         {showIcon && <Icon name="clock-nine" size={'sm'} className={styles.icon} />}
 
@@ -99,20 +118,22 @@ export const TimeRangeInput = ({
         )}
       </button>
       {isOpen && (
-        <ClickOutsideWrapper includeButtonPress={false} onClick={onClose}>
-          <TimePickerContent
-            timeZone={timeZone}
-            value={isValidTimeRange(value) ? value : getDefaultTimeRange()}
-            onChange={onRangeChange}
-            quickOptions={quickOptions}
-            onChangeTimeZone={onChangeTimeZone}
-            className={styles.content}
-            hideTimeZone={hideTimeZone}
-            isReversed={isReversed}
-            hideQuickRanges={hideQuickRanges}
-            weekStart={weekStart}
-          />
-        </ClickOutsideWrapper>
+        <FocusScope contain autoFocus restoreFocus>
+          <section className={styles.content} ref={overlayRef} {...overlayProps} {...dialogProps}>
+            <TimePickerContent
+              timeZone={timeZone}
+              value={isValidTimeRange(value) ? value : getDefaultTimeRange()}
+              onChange={onRangeChange}
+              quickOptions={quickOptions}
+              onChangeTimeZone={onChangeTimeZone}
+              className={styles.content}
+              hideTimeZone={hideTimeZone}
+              isReversed={isReversed}
+              hideQuickRanges={hideQuickRanges}
+              weekStart={weekStart}
+            />
+          </section>
+        </FocusScope>
       )}
     </div>
   );