Fix issue with switching build without changing story

src/screens/VisualTests/VisualTests.stories.tsx

@@ -118,6 +118,7 @@ const meta = {
   },
   args: {
     setSelectedBuildInfo: action("setSelectedBuildInfo"),
+    dismissBuildError: action("dismissBuildError"),
     gitInfo: {
       userEmailHash: "xyz987",
       branch: "feature-branch",
@@ -372,10 +373,6 @@ export const StoryAddedInSelectedBuild = {
   },
 } satisfies Story;
 
-/**
- * Although this state doesn't immediately render the captured story (it probably should),
- * it should switch to the lastBuildOnBranch immediately.
- */
 export const StoryAddedInLastBuildOnBranchNotInSelected = {
   args: {
     selectedBuildInfo: { buildId: pendingBuild.id, storyId: meta.args.storyId },
@@ -573,30 +570,29 @@ export const PendingLocalBuildCapturedStory = {
       "https://www.figma.com/file/GFEbCgCVDtbZhngULbw2gP/Visual-testing-in-Storybook?type=design&node-id=2303-374529&t=qjmuGHxoALrVuhvX-0"
     ),
   },
-  // NOTE: this does not current work, the story is auto-selected
-  // play: async ({ canvasElement, args, argsByTarget }) => {
-  //   // We have to wait a moment for the story to render
-  //   const canvas = within(canvasElement);
-  //   await canvas.findAllByText("Latest");
+  play: async ({ canvasElement, args, argsByTarget }) => {
+    // We have to wait a moment for the story to render
+    const canvas = within(canvasElement);
+    await canvas.findAllByText("Latest");
 
-  //   // Ensure we don't switch to the new build, the user has to opt-in
-  //   mock(args.setSelectedBuildInfo!).mock.calls.forEach(([updater]) => {
-  //     const result = typeof updater === "function" ? updater(args.selectedBuildInfo) : updater;
-  //     expect(result).toEqual(args.selectedBuildInfo); // Unchanged
-  //   });
+    // Ensure we don't switch to the new build, the user has to opt-in
+    mock(args.setSelectedBuildInfo!).mock.calls.forEach(([updater]) => {
+      const result = typeof updater === "function" ? updater(args.selectedBuildInfo) : updater;
+      expect(result).toEqual(args.selectedBuildInfo); // Unchanged
+    });
 
-  //   // Now opt in
-  //   const viewLatestSnapshot = (await canvas.findAllByText("View latest snapshot"))[0];
-  //   await userEvent.click(viewLatestSnapshot);
+    // Now opt in
+    const viewLatestSnapshot = (await canvas.findAllByText("View latest snapshot"))[0];
+    await userEvent.click(viewLatestSnapshot);
 
-  //   const graphqlArgs = argsByTarget.graphql?.$graphql as typeof args.$graphql; // We need to type argsByTarget
-  //   await waitFor(() => {
-  //     expect(args.setSelectedBuildInfo).toHaveBeenCalledWith({
-  //       buildId: graphqlArgs?.AddonVisualTestsBuild?.lastBuildOnBranch?.id,
-  //       storyId: meta.args.storyId,
-  //     });
-  //   });
-  // },
+    const graphqlArgs = argsByTarget.graphql?.$graphql as typeof args.$graphql; // We need to type argsByTarget
+    await waitFor(() => {
+      expect(args.setSelectedBuildInfo).toHaveBeenCalledWith({
+        buildId: graphqlArgs?.AddonVisualTestsBuild?.lastBuildOnBranch?.id,
+        storyId: meta.args.storyId,
+      });
+    });
+  },
 } satisfies Story;
 
 /**
@@ -622,9 +618,25 @@ export const PendingCIBuildCapturedStory = {
       "https://www.figma.com/file/GFEbCgCVDtbZhngULbw2gP/Visual-testing-in-Storybook?type=design&node-id=2303-374529&t=qjmuGHxoALrVuhvX-0"
     ),
   },
-  // NOTE: this does not current work, the story is auto-selected
   // Should behave the same as when the local build captures the current story
-  // play: PendingLocalBuildCapturedStory.play,
+  play: PendingLocalBuildCapturedStory.play,
+} satisfies Story;
+
+/**
+ * Now what happens when the user actually clicks to change story?
+ */
+export const PendingCIBuildCapturedStoryAndUserChangedStory = {
+  args: {
+    selectedBuildInfo: { buildId: pendingBuild.id, storyId: "old--story" },
+    $graphql: PendingLocalBuildCapturedStory.args.$graphql,
+  },
+  parameters: {
+    ...withFigmaDesign(
+      "https://www.figma.com/file/GFEbCgCVDtbZhngULbw2gP/Visual-testing-in-Storybook?type=design&node-id=2303-374529&t=qjmuGHxoALrVuhvX-0"
+    ),
+  },
+  // In this case, we *should* switch to the new build
+  play: EmptyBranchLocalBuildCapturedCurrentStory.play,
 } satisfies Story;
 
 export const Pending = {
@@ -862,8 +874,7 @@ export const CIBuildNewer = {
     },
   },
   // Similarly to a captured story, we should only switch when the user is ready
-  // NOTE: this does not current work, the story is auto-selected
-  // play: PendingLocalBuildCapturedStory.play,
+  play: PendingLocalBuildCapturedStory.play,
 } satisfies Story;
 
 /** The new build is newer than the story build and the git info */

src/screens/VisualTests/VisualTests.tsx

@@ -154,12 +154,22 @@ export const VisualTestsWithoutSelectedBuildId = ({
     !!lastBuildOnBranch &&
     lastBuildOnBranchStoryTests.every(({ status }) => status !== TestStatus.InProgress);
 
+  // We may end up using the lastBuildOnBranch as selected build if we know we are about
+  // to switch to it.
+  const selectedBuildFromData = getFragment(FragmentSelectedBuildFields, data?.selectedBuild);
+
+  const selectedBuildHasStory =
+    selectedBuildFromData &&
+    "testsForStory" in selectedBuildFromData &&
+    selectedBuildFromData.testsForStory?.nodes?.length;
+
   // Before we set the storyInfo, we use the lastBuildOnBranch for story data if it's ready
-  const selectedBuild = getFragment(
-    FragmentSelectedBuildFields,
-    data?.selectedBuild ??
-      (lastBuildOnBranchCompletedStory ? data?.project?.lastBuildOnBranch : undefined)
-  );
+  const selectedBuild = selectedBuildHasStory
+    ? selectedBuildFromData
+    : getFragment(
+        FragmentSelectedBuildFields,
+        lastBuildOnBranchCompletedStory ? data?.project?.lastBuildOnBranch : undefined
+      );
 
   const selectedBuildHasCorrectBranch = selectedBuild?.branch === gitInfo.branch;
   // Currently only used by the sidebar button to show a blue dot ("build outdated")
@@ -195,19 +205,30 @@ export const VisualTestsWithoutSelectedBuildId = ({
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [JSON.stringify(buildStatusUpdate), updateBuildStatus]);
 
+  // Should means "we'd like you to but we'll let you decide"
   const shouldSwitchToLastBuildOnBranch =
     canSwitchToLastBuildOnBranch && lastBuildOnBranchCompletedStory;
 
+  // Force means "we are going to change you to this build"
+  const forceSwitchToLastBuildOnBranch = shouldSwitchToLastBuildOnBranch && !selectedBuildHasStory;
+
   // Ensure we are holding the right story build
   useEffect(() => {
     setSelectedBuildInfo((oldSelectedBuildInfo) =>
       updateSelectedBuildInfo(oldSelectedBuildInfo, {
         shouldSwitchToLastBuildOnBranch,
+        forceSwitchToLastBuildOnBranch,
         lastBuildOnBranchId: lastBuildOnBranch?.id,
         storyId,
       })
     );
-  }, [shouldSwitchToLastBuildOnBranch, lastBuildOnBranch?.id, storyId, setSelectedBuildInfo]);
+  }, [
+    shouldSwitchToLastBuildOnBranch,
+    forceSwitchToLastBuildOnBranch,
+    lastBuildOnBranch?.id,
+    storyId,
+    setSelectedBuildInfo,
+  ]);
 
   const switchToLastBuildOnBranch = useCallback(
     () =>

src/utils/updateSelectedBuildInfo.test.ts

@@ -1,102 +1,103 @@
 import { updateSelectedBuildInfo } from "./updateSelectedBuildInfo";
 
-it("does nothing if there is no next build", () => {
-  expect(
-    updateSelectedBuildInfo(undefined, {
-      shouldSwitchToLastBuildOnBranch: false,
-      lastBuildOnBranchId: undefined,
-      storyId: "storyId",
-    })
-  ).toEqual(undefined);
-});
-
-it("sets the story build from the next build, simple", () => {
-  expect(
-    updateSelectedBuildInfo(undefined, {
-      shouldSwitchToLastBuildOnBranch: true,
-      lastBuildOnBranchId: "lastBuildOnBranchId",
-      storyId: "storyId",
-    })
-  ).toEqual({
-    buildId: "lastBuildOnBranchId",
-    storyId: "storyId",
+describe("with no selected build", () => {
+  it("does nothing if there is no last build on branch", () => {
+    expect(
+      updateSelectedBuildInfo(undefined, {
+        shouldSwitchToLastBuildOnBranch: false,
+        forceSwitchToLastBuildOnBranch: false,
+        lastBuildOnBranchId: undefined,
+        storyId: "storyId",
+      })
+    ).toEqual(undefined);
   });
-});
-
-// We should remain on the "new build" screen until we see a completed story
-it("does not set the story build from the next build, if the next build should not be switched to", () => {
-  expect(
-    updateSelectedBuildInfo(undefined, {
-      shouldSwitchToLastBuildOnBranch: false,
-      lastBuildOnBranchId: "lastBuildOnBranchId",
-      storyId: "storyId",
-    })
-  ).toEqual(undefined);
-});
 
-it("updates the story build from the next build, simple", () => {
-  expect(
-    updateSelectedBuildInfo(
-      { buildId: "oldBuildId", storyId: "storyId" },
-      {
+  it("sets the selected build from the last build on branch, simple", () => {
+    expect(
+      updateSelectedBuildInfo(undefined, {
         shouldSwitchToLastBuildOnBranch: true,
+        forceSwitchToLastBuildOnBranch: false,
         lastBuildOnBranchId: "lastBuildOnBranchId",
         storyId: "storyId",
-      }
-    )
-  ).toEqual({
-    buildId: "lastBuildOnBranchId",
-    storyId: "storyId",
+      })
+    ).toEqual({
+      buildId: "lastBuildOnBranchId",
+      storyId: "storyId",
+    });
   });
-});
 
-it("does not update the story build from the next build, if the next build should not be switched to", () => {
-  expect(
-    updateSelectedBuildInfo(
-      { buildId: "oldBuildId", storyId: "storyId" },
-      {
+  // We should remain on the "new build" screen until we see a completed story
+  it("does not set the selected build from the last build on branch, if the last build on branch should not be switched to", () => {
+    expect(
+      updateSelectedBuildInfo(undefined, {
         shouldSwitchToLastBuildOnBranch: false,
+        forceSwitchToLastBuildOnBranch: false,
         lastBuildOnBranchId: "lastBuildOnBranchId",
         storyId: "storyId",
-      }
-    )
-  ).toEqual({ buildId: "oldBuildId", storyId: "storyId" });
+      })
+    ).toEqual(undefined);
+  });
 });
 
-it("updates the storyId, simple", () => {
-  expect(
-    updateSelectedBuildInfo(
-      {
-        buildId: "lastBuildOnBranchId",
-        storyId: "storyId",
-      },
-      {
-        shouldSwitchToLastBuildOnBranch: true,
-        lastBuildOnBranchId: "lastBuildOnBranchId",
-        storyId: "newStoryId",
-      }
-    )
-  ).toEqual({
-    buildId: "lastBuildOnBranchId",
-    storyId: "newStoryId",
+describe("with a selected build, when not changing story", () => {
+  it("does not update the selected build from the last build on branch, even if we should", () => {
+    expect(
+      updateSelectedBuildInfo(
+        { buildId: "oldBuildId", storyId: "storyId" },
+        {
+          shouldSwitchToLastBuildOnBranch: true,
+          forceSwitchToLastBuildOnBranch: false,
+          lastBuildOnBranchId: "lastBuildOnBranchId",
+          storyId: "storyId",
+        }
+      )
+    ).toEqual({ buildId: "oldBuildId", storyId: "storyId" });
+  });
+
+  it("does update the selected build from the last build on branch, if forced", () => {
+    expect(
+      updateSelectedBuildInfo(
+        { buildId: "oldBuildId", storyId: "storyId" },
+        {
+          shouldSwitchToLastBuildOnBranch: true,
+          forceSwitchToLastBuildOnBranch: true,
+          lastBuildOnBranchId: "lastBuildOnBranchId",
+          storyId: "storyId",
+        }
+      )
+    ).toEqual({ buildId: "lastBuildOnBranchId", storyId: "storyId" });
   });
 });
 
-it("updates the storyId, keeping the old build if the next build should not be switched to", () => {
-  expect(
-    updateSelectedBuildInfo(
-      {
-        buildId: "oldBuildId",
-        storyId: "storyId",
-      },
-      {
-        shouldSwitchToLastBuildOnBranch: false,
-        lastBuildOnBranchId: "lastBuildOnBranchId",
-        storyId: "newStoryId",
-      }
-    )
-  ).toEqual({
-    buildId: "oldBuildId",
-    storyId: "newStoryId",
+describe("with a selected build, when changing story", () => {
+  it("updates the selected build from the last build on branch, simple", () => {
+    expect(
+      updateSelectedBuildInfo(
+        { buildId: "oldBuildId", storyId: "storyId" },
+        {
+          shouldSwitchToLastBuildOnBranch: true,
+          forceSwitchToLastBuildOnBranch: false,
+          lastBuildOnBranchId: "lastBuildOnBranchId",
+          storyId: "newStoryId",
+        }
+      )
+    ).toEqual({
+      buildId: "lastBuildOnBranchId",
+      storyId: "newStoryId",
+    });
+  });
+
+  it("does not update the selected build from the last build on branch, if the last build on branch should not be switched to", () => {
+    expect(
+      updateSelectedBuildInfo(
+        { buildId: "oldBuildId", storyId: "storyId" },
+        {
+          shouldSwitchToLastBuildOnBranch: false,
+          forceSwitchToLastBuildOnBranch: false,
+          lastBuildOnBranchId: "lastBuildOnBranchId",
+          storyId: "newStoryId",
+        }
+      )
+    ).toEqual({ buildId: "oldBuildId", storyId: "newStoryId" });
   });
 });

src/utils/updateSelectedBuildInfo.ts

@@ -7,14 +7,21 @@ export function updateSelectedBuildInfo(
   oldSelectedBuildInfo: SelectedBuildInfo | undefined,
   {
     shouldSwitchToLastBuildOnBranch,
+    forceSwitchToLastBuildOnBranch,
     lastBuildOnBranchId,
     storyId,
   }: {
     shouldSwitchToLastBuildOnBranch: boolean;
+    forceSwitchToLastBuildOnBranch: boolean;
     lastBuildOnBranchId?: string;
     storyId: string;
   }
 ) {
+  // Never touch the selected build if we don't change story
+  if (oldSelectedBuildInfo?.storyId === storyId && !forceSwitchToLastBuildOnBranch) {
+    return oldSelectedBuildInfo;
+  }
+
   if (!shouldSwitchToLastBuildOnBranch) {
     if (!oldSelectedBuildInfo) return undefined;