coder · aslilac · Aug 22, 2023 · Aug 18, 2023 · Aug 18, 2023 · Aug 18, 2023
diff --git a/site/src/pages/SetupPage/SetupPage.test.tsx b/site/src/pages/SetupPage/SetupPage.test.tsx
@@ -1,7 +1,8 @@
 import { fireEvent, screen, waitFor } from "@testing-library/react"
 import userEvent from "@testing-library/user-event"
 import { rest } from "msw"
-import { render } from "testHelpers/renderHelpers"
+import { createMemoryRouter } from "react-router-dom"
+import { render, renderWithRouter } from "testHelpers/renderHelpers"
 import { server } from "testHelpers/server"
 import { SetupPage } from "./SetupPage"
 import { Language as PageViewLanguage } from "./SetupPageView"
@@ -35,6 +36,12 @@ describe("Setup Page", () => {
       rest.get("/api/v2/users/me", (req, res, ctx) => {
         return res(ctx.status(401), ctx.json({ message: "no user here" }))
       }),
+      rest.get("/api/v2/users/first", (req, res, ctx) => {
+        return res(
+          ctx.status(404),
+          ctx.json({ message: "no first user has been created" }),
+        )
+      }),
     )
   })
 
@@ -63,23 +70,109 @@ describe("Setup Page", () => {
         )
       }),
     )
+
     render(<SetupPage />)
     await fillForm()
     const errorMessage = await screen.findByText(fieldErrorMessage)
     expect(errorMessage).toBeDefined()
   })
 
-  it("redirects to workspaces page when success", async () => {
+  it("redirects to the app when setup is successful", async () => {
+    let userHasBeenCreated = false
+
+    server.use(
+      rest.get("/api/v2/users/me", (req, res, ctx) => {
+        if (!userHasBeenCreated) {
+          return res(ctx.status(401), ctx.json({ message: "no user here" }))
+        }
+        return res(ctx.status(200), ctx.json(MockUser))
+      }),
+      rest.get("/api/v2/users/first", (req, res, ctx) => {
+        if (!userHasBeenCreated) {
+          return res(
+            ctx.status(404),
+            ctx.json({ message: "no first user has been created" }),
+          )
+        }
+        return res(
+          ctx.status(200),
+          ctx.json({ message: "hooray, someone exists!" }),
+        )
+      }),
+      rest.post("/api/v2/users/first", (req, res, ctx) => {
+        userHasBeenCreated = true
+        return res(
+          ctx.status(200),
+          ctx.json({ data: "user setup was successful!" }),
+        )
+      }),
+    )
+
     render(<SetupPage />)
+    await fillForm()
+    await waitFor(() => expect(window.location).toBeAt("/"))
+  })
+
+  it("redirects to login if setup has already completed", async () => {
+    // simulates setup having already been completed
+    server.use(
+      rest.get("/api/v2/users/first", (req, res, ctx) => {
+        return res(
+          ctx.status(200),
+          ctx.json({ message: "hooray, someone exists!" }),
+        )
+      }),
+    )
+
+    renderWithRouter(
+      createMemoryRouter(
+        [
+          {
+            path: "/setup",
+            element: <SetupPage />,
+          },
+          {
+            path: "/login",
+            element: <h1>Login</h1>,
+          },
+        ],
+        { initialEntries: ["/setup"] },
+      ),
+    )
 
+    await screen.findByText("Login")
+  })
+
+  it("redirects to the app when already logged in", async () => {
     // simulates the user will be authenticated
     server.use(
       rest.get("/api/v2/users/me", (req, res, ctx) => {
         return res(ctx.status(200), ctx.json(MockUser))
       }),
+      rest.get("/api/v2/users/first", (req, res, ctx) => {
+        return res(
+          ctx.status(200),
+          ctx.json({ message: "hooray, someone exists!" }),
+        )
+      }),
     )
 
-    await fillForm()
-    await waitFor(() => expect(window.location).toBeAt("/workspaces"))
+    renderWithRouter(
+      createMemoryRouter(
+        [
+          {
+            path: "/setup",
+            element: <SetupPage />,
+          },
+          {
+            path: "/",
+            element: <h1>Workspaces</h1>,
+          },
+        ],
+        { initialEntries: ["/setup"] },
+      ),
+    )
+
+    await screen.findByText("Workspaces")
   })
 })
diff --git a/site/src/pages/SetupPage/SetupPage.tsx b/site/src/pages/SetupPage/SetupPage.tsx
@@ -1,10 +1,11 @@
 import { useMachine } from "@xstate/react"
 import { useAuth } from "components/AuthProvider/AuthProvider"
-import { FC, useEffect } from "react"
+import { FC } from "react"
 import { Helmet } from "react-helmet-async"
 import { pageTitle } from "utils/page"
 import { setupMachine } from "xServices/setup/setupXService"
 import { SetupPageView } from "./SetupPageView"
+import { Navigate } from "react-router-dom"
 
 export const SetupPage: FC = () => {
   const [authState, authSend] = useAuth()
@@ -24,11 +25,20 @@ export const SetupPage: FC = () => {
   })
   const { error } = setupState.context
 
-  useEffect(() => {
-    if (authState.matches("signedIn")) {
-      window.location.assign("/workspaces")
-    }
-  }, [authState])
+  const userIsSignedIn = authState.matches("signedIn")
+  const setupIsComplete =
+    !authState.matches("loadingInitialAuthData") &&
+    !authState.matches("configuringTheFirstUser")
+
+  // If the user is logged in, navigate to the app
+  if (userIsSignedIn) {
+    return <Navigate to="/" state={{ isRedirect: true }} />
+  }
+
+  // If we've already completed setup, navigate to the login page
+  if (setupIsComplete) {
+    return <Navigate to="/login" state={{ isRedirect: true }} />
+  }
 
   return (
     <>