Fix notebook editor shows the remove button on the data step (#29575)

* Add basic tests to `NotebookStep`

* Fix irreversible steps have "remove" button

* Use `aria-label` instead of a test ID

* Use `IconButtonWrapper`

e2e/test/scenarios/custom-column/reproductions/19745-cc-nested-query-remove-expressions.cy.spec.js

@@ -97,7 +97,7 @@ function removeExpression(name) {
 
 function removeAllExpressions() {
   getNotebookStep("expression", { stage: 1 }).within(() => {
-    cy.findByTestId("remove-step").click({ force: true });
+    cy.findByLabelText("Remove step").click({ force: true });
   });
 }
 

e2e/test/scenarios/question/reproductions/17514-ui-overlay.cy.spec.js

@@ -172,7 +172,7 @@ function removeJoinedTable() {
   cy.findAllByText("Join data")
     .first()
     .parent()
-    .findByTestId("remove-step")
+    .findByLabelText("Remove step")
     .click({ force: true });
 }
 

frontend/src/metabase/query_builder/components/notebook/NotebookStep/NotebookStep.tsx

@@ -5,6 +5,7 @@ import { color as c } from "metabase/lib/colors";
 import { useToggle } from "metabase/hooks/use-toggle";
 
 import Icon from "metabase/components/Icon";
+import IconButtonWrapper from "metabase/components/IconButtonWrapper";
 import ExpandingContent from "metabase/components/ExpandingContent";
 
 import type Question from "metabase-lib/Question";
@@ -104,6 +105,7 @@ function NotebookStep({
   const color = getColor();
   const canPreview = step?.previewQuery?.isValid?.();
   const hasPreviewButton = !isPreviewOpen && canPreview;
+  const canRevert = typeof step.revert === "function";
 
   return (
     <ExpandingContent isInitiallyOpen={!isLastOpened} isOpen>
@@ -113,13 +115,18 @@ function NotebookStep({
       >
         <StepHeader color={color}>
           {title}
-          <Icon
-            name="close"
-            className="ml-auto cursor-pointer text-light text-medium-hover hover-child"
-            tooltip={t`Remove`}
-            onClick={handleClickRevert}
-            data-testid="remove-step"
-          />
+          {canRevert && (
+            <IconButtonWrapper
+              className="ml-auto text-light text-medium-hover hover-child"
+              onClick={handleClickRevert}
+            >
+              <Icon
+                name="close"
+                tooltip={t`Remove`}
+                aria-label={t`Remove step`}
+              />
+            </IconButtonWrapper>
+          )}
         </StepHeader>
 
         {NotebookStepComponent && (

frontend/src/metabase/query_builder/components/notebook/NotebookStep/NotebookStep.unit.spec.tsx

+import React from "react";
+import { renderWithProviders, screen } from "__support__/ui";
+import {
+  setupDatabasesEndpoints,
+  setupSearchEndpoints,
+} from "__support__/server-mocks";
+
+import { createSampleDatabase } from "metabase-types/api/mocks/presets";
+
+import { getSavedStructuredQuestion } from "metabase-lib/mocks";
+import type Question from "metabase-lib/Question";
+import type StructuredQuery from "metabase-lib/queries/StructuredQuery";
+
+import {
+  NotebookStep as INotebookStep,
+  NotebookStepType,
+} from "../lib/steps.types";
+import NotebookStep from "./NotebookStep";
+
+type SetupOpts = {
+  step?: INotebookStep;
+  question?: Question;
+};
+
+const DEFAULT_QUESTION = getSavedStructuredQuestion();
+const DEFAULT_QUERY = DEFAULT_QUESTION.query() as StructuredQuery;
+
+function createNotebookStep(opts = {}): INotebookStep {
+  return {
+    id: "test-step",
+    type: "data",
+    stageIndex: 0,
+    itemIndex: 0,
+    query: DEFAULT_QUERY,
+    valid: true,
+    active: true,
+    visible: true,
+    actions: [],
+    previewQuery: null,
+    next: null,
+    previous: null,
+    revert: jest.fn(),
+    clean: jest.fn(),
+    update: jest.fn(),
+    ...opts,
+  };
+}
+
+function setup({
+  step = createNotebookStep(),
+  question = DEFAULT_QUESTION,
+}: SetupOpts = {}) {
+  const openStep = jest.fn();
+  const updateQuery = jest.fn();
+
+  setupDatabasesEndpoints([createSampleDatabase()]);
+  setupSearchEndpoints([]);
+
+  renderWithProviders(
+    <NotebookStep
+      step={step}
+      sourceQuestion={question}
+      isLastStep={false}
+      isLastOpened={false}
+      openStep={openStep}
+      updateQuery={updateQuery}
+    />,
+  );
+
+  return {
+    openStep,
+    updateQuery,
+  };
+}
+
+const STEP_TYPES: NotebookStepType[] = [
+  "data",
+  "join",
+  "expression",
+  "filter",
+  "summarize",
+  "aggregate",
+  "breakout",
+  "sort",
+  "limit",
+];
+
+describe("NotebookStep", () => {
+  test.each(STEP_TYPES)(`renders a %s step correctly`, type => {
+    const step = createNotebookStep({ type });
+    const testId = `step-${type}-${step.stageIndex}-${step.itemIndex}`;
+    setup({ step });
+
+    expect(screen.getByTestId(testId)).toBeInTheDocument();
+    expect(
+      screen.getByRole("button", { name: "Remove step" }),
+    ).toBeInTheDocument();
+  });
+
+  it("doesn't render the remove button if step revert isn't implemented", () => {
+    const step = createNotebookStep({ type: "data", revert: null });
+    setup({ step });
+
+    expect(
+      screen.queryByRole("button", { name: "Remove step" }),
+    ).not.toBeInTheDocument();
+  });
+});