From 3047ba70e119cc1d953bb420b4c1a88611632e00 Mon Sep 17 00:00:00 2001
From: Anton Kulyk <kuliks.anton@gmail.com>
Date: Thu, 18 Nov 2021 15:31:49 +0200
Subject: [PATCH] Adding collection items from collection page (#19009)

---
 frontend/src/metabase-lib/lib/Question.js     |   3 +
 .../CollectionHeader/CollectionHeader.jsx     |  26 +---
 .../CollectionHeader.unit.spec.js             |   8 +-
 .../components/NewCollectionItemMenu.jsx      |  42 ++++++
 .../metabase/containers/CollectionName.jsx    |   6 +-
 .../query_builder/components/DataSelector.jsx | 133 +++++++++++++++++-
 .../components/DataSelector.styled.jsx        |  29 ++++
 .../components/notebook/steps/DataStep.jsx    |  15 +-
 frontend/test/__support__/e2e/cypress.js      |   1 +
 .../e2e/helpers/e2e-collection-helpers.js     |  20 +++
 .../collections/collection-types.cy.spec.js   |  13 +-
 .../collections/permissions.cy.spec.js        |   8 +-
 .../personal-collections.cy.spec.js           |  19 ++-
 .../scenarios/question/datasets.cy.spec.js    |  86 ++++++++++-
 14 files changed, 363 insertions(+), 46 deletions(-)
 create mode 100644 frontend/src/metabase/collections/components/NewCollectionItemMenu.jsx
 create mode 100644 frontend/test/__support__/e2e/helpers/e2e-collection-helpers.js

diff --git a/frontend/src/metabase-lib/lib/Question.js b/frontend/src/metabase-lib/lib/Question.js
index 0b60aa983ae..117d7e232f8 100644
--- a/frontend/src/metabase-lib/lib/Question.js
+++ b/frontend/src/metabase-lib/lib/Question.js
@@ -139,6 +139,7 @@ export default class Question {
   static create({
     databaseId,
     tableId,
+    collectionId,
     metadata,
     parameterValues,
     type = "query",
@@ -161,6 +162,7 @@ export default class Question {
   } = {}) {
     let card: CardObject = {
       name,
+      collection_id: collectionId,
       display,
       visualization_settings,
       dataset_query,
@@ -1073,6 +1075,7 @@ export default class Question {
     const cardCopy = {
       name: this._card.name,
       description: this._card.description,
+      collection_id: this._card.collection_id,
       dataset_query: query.datasetQuery(),
       display: this._card.display,
       parameters: this._card.parameters,
diff --git a/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.jsx b/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.jsx
index c2ffa584e19..912d5da6e91 100644
--- a/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.jsx
+++ b/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.jsx
@@ -9,7 +9,9 @@ import Icon, { IconWrapper } from "metabase/components/Icon";
 import Link from "metabase/components/Link";
 import PageHeading from "metabase/components/type/PageHeading";
 import Tooltip from "metabase/components/Tooltip";
+
 import CollectionEditMenu from "metabase/collections/components/CollectionEditMenu";
+import NewCollectionItemMenu from "metabase/collections/components/NewCollectionItemMenu";
 
 import { PLUGIN_COLLECTION_COMPONENTS } from "metabase/plugins";
 
@@ -83,30 +85,12 @@ function EditMenu({
   ) : null;
 }
 
-function CreateCollectionLink({
-  collection,
-  collectionId,
-  hasWritePermission,
-}) {
-  const tooltip = t`New collection`;
-  const link = Urls.newCollection(collectionId);
-
-  return hasWritePermission ? (
-    <Tooltip tooltip={tooltip}>
-      <Link to={link}>
-        <IconWrapper>
-          <Icon name="new_folder" />
-        </IconWrapper>
-      </Link>
-    </Tooltip>
-  ) : null;
-}
-
 function Menu(props) {
+  const { hasWritePermission } = props;
   return (
-    <MenuContainer>
+    <MenuContainer data-testid="collection-menu">
+      {hasWritePermission && <NewCollectionItemMenu {...props} />}
       <EditMenu {...props} />
-      <CreateCollectionLink {...props} />
       <PermissionsLink {...props} />
     </MenuContainer>
   );
diff --git a/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.unit.spec.js b/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.unit.spec.js
index 77341dbadcd..278c9f49323 100644
--- a/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.unit.spec.js
+++ b/frontend/src/metabase/collections/components/CollectionHeader/CollectionHeader.unit.spec.js
@@ -78,8 +78,8 @@ describe("permissions link", () => {
   });
 });
 
-describe("link to edit collection", () => {
-  const ariaLabel = "pencil icon";
+describe("link to add new collection items", () => {
+  const ariaLabel = "add icon";
 
   describe("should not be displayed", () => {
     it("when no detail is passed in the collection to determine if user can change collection", () => {
@@ -104,8 +104,8 @@ describe("link to edit collection", () => {
   });
 });
 
-describe("link to create a new collection", () => {
-  const ariaLabel = "new_folder icon";
+describe("link to add new collection items", () => {
+  const ariaLabel = "add icon";
 
   describe("should not be displayed", () => {
     it("if user is not allowed to change collection", () => {
diff --git a/frontend/src/metabase/collections/components/NewCollectionItemMenu.jsx b/frontend/src/metabase/collections/components/NewCollectionItemMenu.jsx
new file mode 100644
index 00000000000..772e54fe7b7
--- /dev/null
+++ b/frontend/src/metabase/collections/components/NewCollectionItemMenu.jsx
@@ -0,0 +1,42 @@
+import React from "react";
+import PropTypes from "prop-types";
+import { t } from "ttag";
+
+import * as Urls from "metabase/lib/urls";
+
+import EntityMenu from "metabase/components/EntityMenu";
+import { ANALYTICS_CONTEXT } from "metabase/collections/constants";
+
+const propTypes = {
+  collection: PropTypes.func,
+  list: PropTypes.arrayOf(PropTypes.object),
+};
+
+function NewCollectionItemMenu({ collection }) {
+  const items = [
+    {
+      icon: "insight",
+      title: t`Question`,
+      link: Urls.newQuestion({ mode: "notebook", collectionId: collection.id }),
+      event: `${ANALYTICS_CONTEXT};New Item Menu;Question Click`,
+    },
+    {
+      icon: "dashboard",
+      title: t`Dashboard`,
+      link: Urls.newDashboard(collection.id),
+      event: `${ANALYTICS_CONTEXT};New Item Menu;Dashboard Click`,
+    },
+    {
+      icon: "folder",
+      title: t`Collection`,
+      link: Urls.newCollection(collection.id),
+      event: `${ANALYTICS_CONTEXT};New Item Menu;Collection Click`,
+    },
+  ];
+
+  return <EntityMenu items={items} triggerIcon="add" tooltip={t`New…`} />;
+}
+
+NewCollectionItemMenu.propTypes = propTypes;
+
+export default NewCollectionItemMenu;
diff --git a/frontend/src/metabase/containers/CollectionName.jsx b/frontend/src/metabase/containers/CollectionName.jsx
index 756b26f0f4e..f43160ec82b 100644
--- a/frontend/src/metabase/containers/CollectionName.jsx
+++ b/frontend/src/metabase/containers/CollectionName.jsx
@@ -4,10 +4,10 @@ import React from "react";
 import Collection, { ROOT_COLLECTION } from "metabase/entities/collections";
 
 const CollectionName = ({ id }) => {
-  if (id === undefined || isNaN(id)) {
-    return null;
-  } else if (id === "root" || id === null) {
+  if (id === "root" || id === null) {
     return <span>{ROOT_COLLECTION.name}</span>;
+  } else if (id === undefined || isNaN(id)) {
+    return null;
   } else {
     return <Collection.Name id={id} />;
   }
diff --git a/frontend/src/metabase/query_builder/components/DataSelector.jsx b/frontend/src/metabase/query_builder/components/DataSelector.jsx
index df1522cbea4..070166c8aa9 100644
--- a/frontend/src/metabase/query_builder/components/DataSelector.jsx
+++ b/frontend/src/metabase/query_builder/components/DataSelector.jsx
@@ -1,5 +1,5 @@
 /* eslint-disable react/prop-types */
-import React, { Component } from "react";
+import React, { Component, useEffect, useState } from "react";
 import { connect } from "react-redux";
 import PropTypes from "prop-types";
 import { t } from "ttag";
@@ -8,6 +8,7 @@ import _ from "underscore";
 
 import {
   isVirtualCardId,
+  getQuestionVirtualTableId,
   SAVED_QUESTIONS_VIRTUAL_DB_ID,
 } from "metabase/lib/saved-questions";
 
@@ -25,6 +26,9 @@ import Databases from "metabase/entities/databases";
 import Schemas from "metabase/entities/schemas";
 import Tables from "metabase/entities/tables";
 import Search from "metabase/entities/search";
+
+import { PLUGIN_MODERATION } from "metabase/plugins";
+
 import {
   SearchResults,
   convertSearchResultToTableLikeItem,
@@ -38,6 +42,8 @@ import {
   DataBucketListItem,
   PickerSpinner,
   RawDataBackButton,
+  CollectionDatasetSelectList,
+  CollectionDatasetAllDataLink,
 } from "./DataSelector.styled";
 import "./DataSelector.css";
 
@@ -60,6 +66,10 @@ const TABLE_STEP = "TABLE";
 // chooses a table field (table has already been selected)
 const FIELD_STEP = "FIELD";
 
+// allows to choose one of collection's dataset (requires collectionId prop)
+// is used while adding a question with the "+" button on collection page
+const COLLECTION_DATASET_STEP = "COLLECTION_DATASET";
+
 export const DataSourceSelector = props => (
   <DataSelector
     steps={[DATA_BUCKET_STEP, DATABASE_STEP, SCHEMA_STEP, TABLE_STEP]}
@@ -69,6 +79,31 @@ export const DataSourceSelector = props => (
   />
 );
 
+export const CollectionDatasetOrDataSourceSelector = ({
+  hasCollectionDatasetsStep,
+  ...props
+}) => {
+  const [collectionDatasetsShown, setCollectionDatasetsShown] = useState(
+    !!hasCollectionDatasetsStep,
+  );
+
+  const steps = collectionDatasetsShown
+    ? [COLLECTION_DATASET_STEP]
+    : [DATA_BUCKET_STEP, DATABASE_STEP, SCHEMA_STEP, TABLE_STEP];
+
+  return (
+    <DataSelector
+      steps={steps}
+      combineDatabaseSchemaSteps
+      getTriggerElementContent={TableTriggerContent}
+      {...props}
+      onCloseCollectionDatasets={() => {
+        setCollectionDatasetsShown(false);
+      }}
+    />
+  );
+};
+
 export const DatabaseDataSelector = props => (
   <DataSelector
     steps={[DATABASE_STEP]}
@@ -806,6 +841,15 @@ export class UnconnectedDataSelector extends Component {
     };
 
     switch (this.state.activeStep) {
+      case COLLECTION_DATASET_STEP:
+        return (
+          <CollectionDatasetPicker
+            {...props}
+            collectionId={this.props.collectionId}
+            handleCollectionDatasetSelect={this.handleCollectionDatasetSelect}
+            onSeeAllData={this.handleCollectionDatasetsPickerClose}
+          />
+        );
       case DATA_BUCKET_STEP:
         return <DataBucketPicker {...props} />;
       case DATABASE_STEP:
@@ -863,6 +907,22 @@ export class UnconnectedDataSelector extends Component {
       searchText,
     });
 
+  handleCollectionDatasetSelect = async dataset => {
+    const tableId = getQuestionVirtualTableId(dataset);
+    await this.props.fetchFields(tableId);
+    if (this.props.setSourceTableFn) {
+      this.props.setSourceTableFn(tableId);
+    }
+    this.popover.current.toggle();
+    this.props.onCloseCollectionDatasets();
+    this.switchToStep(TABLE_STEP);
+  };
+
+  handleCollectionDatasetsPickerClose = () => {
+    this.props.onCloseCollectionDatasets();
+    this.switchToStep(this.hasDatasets() ? DATA_BUCKET_STEP : DATABASE_STEP);
+  };
+
   handleSearchItemSelect = async item => {
     const table = convertSearchResultToTableLikeItem(item);
     await this.props.fetchFields(table.id);
@@ -987,6 +1047,77 @@ export class UnconnectedDataSelector extends Component {
   }
 }
 
+const CollectionDatasetPicker = ({
+  collectionId,
+  handleCollectionDatasetSelect,
+  onSeeAllData,
+}) => {
+  return (
+    <Search.ListLoader
+      query={{
+        collection: collectionId,
+        models: ["dataset"],
+      }}
+      loadingAndErrorWrapper={false}
+    >
+      {({ list: datasets }) => (
+        <CollectionDatasetList
+          datasets={datasets}
+          onSelect={handleCollectionDatasetSelect}
+          onSeeAllData={onSeeAllData}
+        />
+      )}
+    </Search.ListLoader>
+  );
+};
+
+function CollectionDatasetList({ datasets, onSelect, onSeeAllData }) {
+  useEffect(() => {
+    if (datasets?.length === 0) {
+      onSeeAllData();
+    } else if (datasets?.length === 1) {
+      onSelect(datasets[0]);
+    }
+  }, [datasets, onSelect, onSeeAllData]);
+
+  // If there are no datasets, in a collection, we just switch to the normal picker
+  // If there is exactly one dataset, we select it and close the picker
+  // The loading indicator is still shown for both cases to prevent flickering
+  // Example: spinner > one dataset shown > it gets selected > the selector closes, everything flickers
+  if (!datasets || datasets.length === 0 || datasets.length === 1) {
+    return <LoadingAndErrorWrapper loading />;
+  }
+
+  return (
+    <CollectionDatasetSelectList>
+      {datasets.map(dataset => {
+        return (
+          <CollectionDatasetSelectList.Item
+            key={dataset.id}
+            name={dataset.name}
+            onSelect={() => onSelect(dataset)}
+            size="small"
+            icon={{ name: "dataset", size: 16 }}
+            rightIcon={PLUGIN_MODERATION.getStatusIcon(
+              dataset.moderated_status,
+            )}
+          />
+        );
+      })}
+      <CollectionDatasetAllDataLink
+        key="all-data"
+        onSelect={onSeeAllData}
+        as={props => <li {...props} />}
+      >
+        <CollectionDatasetAllDataLink.Content>
+          {t`All data`}
+          <Icon key="icon" name="chevronright" size={12} />
+        </CollectionDatasetAllDataLink.Content>
+      </CollectionDatasetAllDataLink>
+    </CollectionDatasetSelectList>
+  );
+}
+
 const DataBucketPicker = ({ onChangeDataBucket }) => {
   const BUCKETS = [
     {
diff --git a/frontend/src/metabase/query_builder/components/DataSelector.styled.jsx b/frontend/src/metabase/query_builder/components/DataSelector.styled.jsx
index 7f4ae3ed754..c9e6822af03 100644
--- a/frontend/src/metabase/query_builder/components/DataSelector.styled.jsx
+++ b/frontend/src/metabase/query_builder/components/DataSelector.styled.jsx
@@ -119,3 +119,32 @@ export function DataBucketListItem(props) {
     </DataBucketListItemContainer>
   );
 }
+
+export const CollectionDatasetSelectList = styled(SelectList)`
+  width: 300px;
+  max-width: 300px;
+  padding: 0.5rem;
+`;
+
+CollectionDatasetSelectList.Item = SelectList.Item;
+
+export const CollectionDatasetAllDataLink = styled(SelectList.BaseItem)`
+  padding: 0.5rem;
+
+  color: ${color("text-light")};
+  font-weight: bold;
+  cursor: pointer;
+
+  :hover {
+    color: ${color("brand")};
+  }
+`;
+
+CollectionDatasetAllDataLink.Content = styled.span`
+  display: flex;
+  align-items: center;
+
+  .Icon {
+    margin-left: ${space(0)};
+  }
+`;
diff --git a/frontend/src/metabase/query_builder/components/notebook/steps/DataStep.jsx b/frontend/src/metabase/query_builder/components/notebook/steps/DataStep.jsx
index 4d310237d57..0ec5fab8036 100644
--- a/frontend/src/metabase/query_builder/components/notebook/steps/DataStep.jsx
+++ b/frontend/src/metabase/query_builder/components/notebook/steps/DataStep.jsx
@@ -3,7 +3,7 @@ import React from "react";
 import { connect } from "react-redux";
 import { t } from "ttag";
 
-import { DataSourceSelector } from "metabase/query_builder/components/DataSelector";
+import { CollectionDatasetOrDataSourceSelector } from "metabase/query_builder/components/DataSelector";
 import { getDatabasesList } from "metabase/query_builder/selectors";
 
 import { NotebookCell, NotebookCellItem } from "../NotebookCell";
@@ -15,8 +15,17 @@ import {
 import FieldsPicker from "./FieldsPicker";
 
 function DataStep({ color, query, updateQuery }) {
+  const question = query.question();
   const table = query.table();
   const canSelectTableColumns = table && query.isRaw();
+
+  const hasCollectionDatasetsStep =
+    question &&
+    !question.isSaved() &&
+    !question.databaseId() &&
+    !question.tableId() &&
+    question.collectionId() !== undefined;
+
   return (
     <NotebookCell color={color}>
       <NotebookCellItem
@@ -36,8 +45,10 @@ function DataStep({ color, query, updateQuery }) {
         rightContainerStyle={FIELDS_PICKER_STYLES.notebookRightItemContainer}
         data-testid="data-step-cell"
       >
-        <DataSourceSelector
+        <CollectionDatasetOrDataSourceSelector
           hasTableSearch
+          collectionId={question.collectionId()}
+          hasCollectionDatasetsStep={hasCollectionDatasetsStep}
           databaseQuery={{ saved: true }}
           selectedDatabaseId={query.databaseId()}
           selectedTableId={query.tableId()}
diff --git a/frontend/test/__support__/e2e/cypress.js b/frontend/test/__support__/e2e/cypress.js
index fd196cba40a..6c77c0a79fb 100644
--- a/frontend/test/__support__/e2e/cypress.js
+++ b/frontend/test/__support__/e2e/cypress.js
@@ -20,6 +20,7 @@ export * from "./helpers/e2e-mock-app-settings-helpers";
 export * from "./helpers/e2e-notebook-helpers";
 export * from "./helpers/e2e-assertion-helpers";
 export * from "./helpers/e2e-cloud-helpers";
+export * from "./helpers/e2e-collection-helpers";
 export * from "./helpers/e2e-data-model-helpers";
 export * from "./helpers/e2e-misc-helpers";
 export * from "./helpers/e2e-email-helpers";
diff --git a/frontend/test/__support__/e2e/helpers/e2e-collection-helpers.js b/frontend/test/__support__/e2e/helpers/e2e-collection-helpers.js
new file mode 100644
index 00000000000..cee54a659ff
--- /dev/null
+++ b/frontend/test/__support__/e2e/helpers/e2e-collection-helpers.js
@@ -0,0 +1,20 @@
+import { popover } from "__support__/e2e/cypress";
+
+export function assertCanAddItemsToCollection() {
+  cy.findByTestId("collection-menu").within(() => {
+    cy.icon("add");
+  });
+}
+
+/**
+ * Clicks the "+" icon on the collection page and selects one of the menu options
+ * @param {"question" | "dashboard" | "collection"} type
+ */
+export function openNewCollectionItemFlowFor(type) {
+  cy.findByTestId("collection-menu").within(() => {
+    cy.icon("add").click();
+  });
+  popover()
+    .findByText(new RegExp(type, "i"))
+    .click();
+}
diff --git a/frontend/test/metabase/scenarios/collections/collection-types.cy.spec.js b/frontend/test/metabase/scenarios/collections/collection-types.cy.spec.js
index fbb51f46ee7..20fca7dcd5e 100644
--- a/frontend/test/metabase/scenarios/collections/collection-types.cy.spec.js
+++ b/frontend/test/metabase/scenarios/collections/collection-types.cy.spec.js
@@ -4,6 +4,7 @@ import {
   sidebar,
   describeWithToken,
   describeWithoutToken,
+  openNewCollectionItemFlowFor,
 } from "__support__/e2e/cypress";
 import { SAMPLE_DATASET } from "__support__/e2e/cypress_sample_dataset";
 
@@ -59,7 +60,7 @@ describeWithToken("collections types", () => {
     cy.findByText("First collection").click();
 
     // Test not visible when creating a new collection
-    cy.icon("new_folder").click();
+    openNewCollectionItemFlowFor("collection");
     modal().within(() => {
       cy.findByText(TREE_UPDATE_REGULAR_MESSAGE).should("not.exist");
       cy.findByText(TREE_UPDATE_OFFICIAL_MESSAGE).should("not.exist");
@@ -110,7 +111,7 @@ describeWithToken("collections types", () => {
 
     openCollection("First collection");
 
-    cy.icon("new_folder").click();
+    openNewCollectionItemFlowFor("collection");
     modal().within(() => {
       assertNoCollectionTypeInput();
       cy.icon("close").click();
@@ -129,7 +130,7 @@ describeWithToken("collections types", () => {
     openCollection("Your personal collection");
     cy.icon("pencil").should("not.exist");
 
-    cy.icon("new_folder").click();
+    openNewCollectionItemFlowFor("collection");
     modal().within(() => {
       assertNoCollectionTypeInput();
       cy.findByLabelText("Name").type("Personal collection child");
@@ -138,7 +139,7 @@ describeWithToken("collections types", () => {
 
     openCollection("Personal collection child");
 
-    cy.icon("new_folder").click();
+    openNewCollectionItemFlowFor("collection");
     modal().within(() => {
       assertNoCollectionTypeInput();
       cy.icon("close").click();
@@ -155,7 +156,7 @@ describeWithoutToken("collection types", () => {
   it("should not be able to manage collection's authority level", () => {
     cy.visit("/collection/root");
 
-    cy.icon("new_folder").click();
+    openNewCollectionItemFlowFor("collection");
     modal().within(() => {
       assertNoCollectionTypeInput();
       cy.icon("close").click();
@@ -301,7 +302,7 @@ function setOfficial(official = true) {
 }
 
 function createAndOpenOfficialCollection({ name }) {
-  cy.icon("new_folder").click();
+  openNewCollectionItemFlowFor("collection");
   modal().within(() => {
     cy.findByLabelText("Name").type(name);
     setOfficial();
diff --git a/frontend/test/metabase/scenarios/collections/permissions.cy.spec.js b/frontend/test/metabase/scenarios/collections/permissions.cy.spec.js
index 6699fe94614..9b2ba117c06 100644
--- a/frontend/test/metabase/scenarios/collections/permissions.cy.spec.js
+++ b/frontend/test/metabase/scenarios/collections/permissions.cy.spec.js
@@ -49,7 +49,9 @@ describe("collection permissions", () => {
                     displaySidebarChildOf("First collection");
                     displaySidebarChildOf("Second collection");
                   });
-                  cy.icon("add").click();
+                  cy.get(".Nav").within(() => {
+                    cy.icon("add").click();
+                  });
                   cy.findByText("New dashboard").click();
                   cy.get(".AdminSelect").findByText("Second collection");
                 });
@@ -242,7 +244,9 @@ describe("collection permissions", () => {
                       .as("title")
                       .contains("Third collection");
                     // Creating new sub-collection at this point shouldn't be possible
-                    cy.icon("new_folder").should("not.exist");
+                    cy.findByTestId("collection-menu").within(() => {
+                      cy.icon("add").should("not.exist");
+                    });
                     // We shouldn't be able to change permissions for an archived collection (the root issue of #12489!)
                     cy.icon("lock").should("not.exist");
                     /**
diff --git a/frontend/test/metabase/scenarios/collections/personal-collections.cy.spec.js b/frontend/test/metabase/scenarios/collections/personal-collections.cy.spec.js
index 0b9dd97e510..60a6c0997c1 100644
--- a/frontend/test/metabase/scenarios/collections/personal-collections.cy.spec.js
+++ b/frontend/test/metabase/scenarios/collections/personal-collections.cy.spec.js
@@ -1,4 +1,11 @@
-import { restore, popover, modal, sidebar } from "__support__/e2e/cypress";
+import {
+  restore,
+  popover,
+  modal,
+  sidebar,
+  assertCanAddItemsToCollection,
+  openNewCollectionItemFlowFor,
+} from "__support__/e2e/cypress";
 import { USERS } from "__support__/e2e/cypress_data";
 
 describe("personal collections", () => {
@@ -31,12 +38,12 @@ describe("personal collections", () => {
     it("shouldn't be able to change permission levels or edit personal collections", () => {
       cy.visit("/collection/root");
       cy.findByText("Your personal collection").click();
-      cy.icon("new_folder");
+      assertCanAddItemsToCollection();
       cy.icon("lock").should("not.exist");
       cy.icon("pencil").should("not.exist");
       // Visit random user's personal collection
       cy.visit("/collection/5");
-      cy.icon("new_folder");
+      assertCanAddItemsToCollection();
       cy.icon("lock").should("not.exist");
       cy.icon("pencil").should("not.exist");
     });
@@ -49,7 +56,7 @@ describe("personal collections", () => {
       sidebar()
         .findByText("Foo")
         .click();
-      cy.icon("new_folder");
+      assertCanAddItemsToCollection();
       cy.icon("pencil");
       cy.icon("lock").should("not.exist");
 
@@ -74,7 +81,7 @@ describe("personal collections", () => {
 
     it.skip("should be able view other users' personal sub-collections (metabase#15339)", () => {
       cy.visit("/collection/5");
-      cy.icon("new_folder").click();
+      openNewCollectionItemFlowFor("collection");
       cy.findByLabelText("Name").type("Foo");
       cy.findByText("Create").click();
       // This repro could possibly change depending on the design decision for this feature implementation
@@ -150,7 +157,7 @@ describe("personal collections", () => {
 });
 
 function addNewCollection(name) {
-  cy.icon("new_folder").click();
+  openNewCollectionItemFlowFor("collection");
   cy.findByLabelText("Name").type(name);
   cy.findByText("Create").click();
 }
diff --git a/frontend/test/metabase/scenarios/question/datasets.cy.spec.js b/frontend/test/metabase/scenarios/question/datasets.cy.spec.js
index adc43d4d315..b3c5159a057 100644
--- a/frontend/test/metabase/scenarios/question/datasets.cy.spec.js
+++ b/frontend/test/metabase/scenarios/question/datasets.cy.spec.js
@@ -1,4 +1,11 @@
-import { restore, modal, popover, visualize } from "__support__/e2e/cypress";
+import {
+  restore,
+  modal,
+  popover,
+  getNotebookStep,
+  openNewCollectionItemFlowFor,
+  visualize,
+} from "__support__/e2e/cypress";
 
 describe("scenarios > datasets", () => {
   beforeEach(() => {
@@ -151,6 +158,83 @@ describe("scenarios > datasets", () => {
       cy.url().should("match", /\/question\/\d+-[a-z0-9-]*$/);
     });
   });
+
+  describe("adding a question to collection from its page", () => {
+    it("should offer to pick one of the collection's datasets by default", () => {
+      cy.request("PUT", "/api/card/1", { dataset: true });
+      cy.request("PUT", "/api/card/2", { dataset: true });
+
+      cy.visit("/collection/root");
+      openNewCollectionItemFlowFor("question");
+
+      cy.findByText("Orders");
+      cy.findByText("Orders, Count");
+      cy.findByText("All data");
+
+      cy.findByText("Datasets").should("not.exist");
+      cy.findByText("Raw Data").should("not.exist");
+      cy.findByText("Saved Questions").should("not.exist");
+      cy.findByText("Sample Dataset").should("not.exist");
+
+      cy.findByText("Orders").click();
+
+      getNotebookStep("data").within(() => {
+        cy.findByText("Orders");
+      });
+
+      cy.button("Visualize");
+    });
+
+    it("should open the default picker after clicking 'All data'", () => {
+      cy.request("PUT", "/api/card/1", { dataset: true });
+      cy.request("PUT", "/api/card/2", { dataset: true });
+
+      cy.visit("/collection/root");
+      openNewCollectionItemFlowFor("question");
+
+      cy.findByText("All data").click({ force: true });
+
+      cy.findByText("Datasets");
+      cy.findByText("Raw Data");
+      cy.findByText("Saved Questions");
+    });
+
+    it("should automatically use the only collection dataset as a data source", () => {
+      cy.request("PUT", "/api/card/2", { dataset: true });
+
+      cy.visit("/collection/root");
+      openNewCollectionItemFlowFor("question");
+
+      getNotebookStep("data").within(() => {
+        cy.findByText("Orders, Count");
+      });
+      cy.button("Visualize");
+    });
+
+    it("should use correct picker if collection has no datasets", () => {
+      cy.request("PUT", "/api/card/1", { dataset: true });
+
+      cy.visit("/collection/9");
+      openNewCollectionItemFlowFor("question");
+
+      cy.findByText("All data").should("not.exist");
+      cy.findByText("Datasets");
+      cy.findByText("Raw Data");
+      cy.findByText("Saved Questions");
+    });
+
+    it("should use correct picker if there are datasets at all", () => {
+      cy.visit("/collection/root");
+      openNewCollectionItemFlowFor("question");
+
+      cy.findByText("All data").should("not.exist");
+      cy.findByText("Datasets").should("not.exist");
+      cy.findByText("Raw Data").should("not.exist");
+
+      cy.findByText("Saved Questions");
+      cy.findByText("Sample Dataset");
+    });
+  });
 });
 
 function openDetailsSidebar() {
-- 
GitLab