coder · BrunoQuaresma · Jan 30, 2024 · Jan 30, 2024 · Jan 30, 2024
diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts
@@ -1,8 +1,11 @@
 import * as API from "api/api";
-import { QueryClient, type QueryOptions } from "react-query";
+import {
+  QueryClient,
+  UseMutationOptions,
+  type QueryOptions,
+} from "react-query";
 import { putWorkspaceExtension } from "api/api";
-import dayjs from "dayjs";
-import { getDeadline, getMaxDeadline, getMinDeadline } from "utils/schedule";
+import { Dayjs } from "dayjs";
 import {
   type WorkspaceBuildParameter,
   type Workspace,
@@ -103,25 +106,12 @@ export function workspaces(config: WorkspacesRequest = {}) {
   } as const satisfies QueryOptions<WorkspacesResponse>;
 }
 
-export const decreaseDeadline = (workspace: Workspace) => {
-  return {
-    mutationFn: (hours: number) => {
-      const proposedDeadline = getDeadline(workspace).subtract(hours, "hours");
-      const newDeadline = dayjs.max(proposedDeadline, getMinDeadline());
-      return putWorkspaceExtension(workspace.id, newDeadline);
-    },
-  };
-};
-
-export const increaseDeadline = (workspace: Workspace) => {
+export const updateDeadline = (
+  workspace: Workspace,
+): UseMutationOptions<void, unknown, Dayjs> => {
   return {
-    mutationFn: (hours: number) => {
-      const proposedDeadline = getDeadline(workspace).add(hours, "hours");
-      const newDeadline = dayjs.min(
-        proposedDeadline,
-        getMaxDeadline(workspace),
-      );
-      return putWorkspaceExtension(workspace.id, newDeadline);
+    mutationFn: (deadline: Dayjs) => {
+      return putWorkspaceExtension(workspace.id, deadline);
     },
   };
 };

diff --git a/site/src/pages/WorkspacePage/WorkspaceScheduleControls.test.tsx b/site/src/pages/WorkspacePage/WorkspaceScheduleControls.test.tsx
@@ -0,0 +1,156 @@
+import { render, screen } from "@testing-library/react";
+import { ThemeProvider } from "contexts/ThemeProvider";
+import { QueryClient, QueryClientProvider, useQuery } from "react-query";
+import { MockWorkspace } from "testHelpers/entities";
+import { WorkspaceScheduleControls } from "./WorkspaceScheduleControls";
+import { workspaceByOwnerAndName } from "api/queries/workspaces";
+import { RouterProvider, createMemoryRouter } from "react-router-dom";
+import userEvent from "@testing-library/user-event";
+import { server } from "testHelpers/server";
+import { rest } from "msw";
+import dayjs from "dayjs";
+import * as API from "api/api";
+import { GlobalSnackbar } from "components/GlobalSnackbar/GlobalSnackbar";
+
+const Wrapper = () => {
+  const { data: workspace } = useQuery(
+    workspaceByOwnerAndName(MockWorkspace.owner_name, MockWorkspace.name),
+  );
+
+  if (!workspace) {
+    return null;
+  }
+
+  return <WorkspaceScheduleControls workspace={workspace} canUpdateSchedule />;
+};
+
+const BASE_DEADLINE = dayjs().add(3, "hour");
+
+const renderScheduleControls = async () => {
+  server.use(
+    rest.get(
+      "/api/v2/users/:username/workspace/:workspaceName",
+      (req, res, ctx) => {
+        return res(
+          ctx.status(200),
+          ctx.json({
+            ...MockWorkspace,
+            latest_build: {
+              ...MockWorkspace.latest_build,
+              deadline: BASE_DEADLINE.toISOString(),
+            },
+          }),
+        );
+      },
+    ),
+  );
+  render(
+    <ThemeProvider>
+      <QueryClientProvider client={new QueryClient()}>
+        <RouterProvider
+          router={createMemoryRouter([{ path: "/", element: <Wrapper /> }])}
+        />
+      </QueryClientProvider>
+      <GlobalSnackbar />
+    </ThemeProvider>,
+  );
+  await screen.findByTestId("schedule-controls");
+  expect(screen.getByText("Stop in 3 hours")).toBeInTheDocument();
+};
+
+test("add 3 hours to deadline", async () => {
+  const user = userEvent.setup();
+  const updateDeadlineSpy = jest
+    .spyOn(API, "putWorkspaceExtension")
+    .mockResolvedValue();
+
+  await renderScheduleControls();
+
+  const addButton = screen.getByRole("button", {
+    name: /add 1 hour to deadline/i,
+  });
+  await user.click(addButton);
+  await user.click(addButton);
+  await user.click(addButton);
+  await screen.findByText(
+    "Workspace shutdown time has been successfully updated.",
+  );
+  expect(screen.getByText("Stop in 6 hours")).toBeInTheDocument();
+
+  // Mocks are used here because the 'usedDeadline' is a dayjs object, which
+  // can't be directly compared.
+  const usedWorkspaceId = updateDeadlineSpy.mock.calls[0][0];
+  const usedDeadline = updateDeadlineSpy.mock.calls[0][1];
+  expect(usedWorkspaceId).toEqual(MockWorkspace.id);
+  expect(usedDeadline.toISOString()).toEqual(
+    BASE_DEADLINE.add(3, "hour").toISOString(),
+  );
+});
+
+test("remove 3 hours to deadline", async () => {
+  const user = userEvent.setup();
+  const updateDeadlineSpy = jest
+    .spyOn(API, "putWorkspaceExtension")
+    .mockResolvedValue();
+
+  await renderScheduleControls();
+
+  const subButton = screen.getByRole("button", {
+    name: /subtract 1 hour from deadline/i,
+  });
+  await user.click(subButton);
+  await user.click(subButton);
+  await screen.findByText(
+    "Workspace shutdown time has been successfully updated.",
+  );
+  expect(screen.getByText("Stop in an hour")).toBeInTheDocument();
+
+  // Mocks are used here because the 'usedDeadline' is a dayjs object, which
+  // can't be directly compared.
+  const usedWorkspaceId = updateDeadlineSpy.mock.calls[0][0];
+  const usedDeadline = updateDeadlineSpy.mock.calls[0][1];
+  expect(usedWorkspaceId).toEqual(MockWorkspace.id);
+  expect(usedDeadline.toISOString()).toEqual(
+    BASE_DEADLINE.subtract(2, "hour").toISOString(),
+  );
+});
+
+test("rollback to previous deadline on error", async () => {
+  const user = userEvent.setup();
+  const initialScheduleMessage = "Stop in 3 hours";
+  jest.spyOn(API, "putWorkspaceExtension").mockRejectedValue({});
+
+  await renderScheduleControls();
+
+  const addButton = screen.getByRole("button", {
+    name: /add 1 hour to deadline/i,
+  });
+  await user.click(addButton);
+  await user.click(addButton);
+  await user.click(addButton);
+  await screen.findByText(
+    "We couldn't update your workspace shutdown time. Please try again.",
+  );
+  // In case of an error, the schedule message should remain unchanged
+  expect(screen.getByText(initialScheduleMessage)).toBeInTheDocument();
+});
+
+test("request is only sent once when clicking multiple times", async () => {
+  const user = userEvent.setup();
+  const updateDeadlineSpy = jest
+    .spyOn(API, "putWorkspaceExtension")
+    .mockResolvedValue();
+
+  await renderScheduleControls();
+
+  const addButton = screen.getByRole("button", {
+    name: /add 1 hour to deadline/i,
+  });
+  await user.click(addButton);
+  await user.click(addButton);
+  await user.click(addButton);
+  await screen.findByText(
+    "Workspace shutdown time has been successfully updated.",
+  );
+  expect(updateDeadlineSpy).toHaveBeenCalledTimes(1);
+});