From 754354dbf9ae3e5dba992e0676e663c512b84265 Mon Sep 17 00:00:00 2001
From: vchske <chskelt@microsoft.com>
Date: Tue, 19 Sep 2023 17:12:41 -0700
Subject: [PATCH] Adds unit tests for vcore quickstart feature code (#1618)

* Safety checkin

* Adding vcoremongo to Main

* Safety checkin

* Adding vcoremongo to Main

* Safety commit

* Safety checkin

* Adding vcoremongo to Main

* Safety commit

* Integrating mongo shell

* Safety checkin

* Adding vcoremongo to Main

* Safety commit

* Integrating mongo shell

* Safety checkin

* Safety commit

* Enable mongo shell in its own tab

* Safety checkin

* Adding vcoremongo to Main

* Safety commit

* Integrating mongo shell

* Safety checkin

* Safety commit

* Safety commit

* Integrating mongo shell

* Safety checkin

* Safety commit

* Enable mongo shell in its own tab

* Adding message

* Integrated mongo shell

* Moving Juno endpoint back to prod

* Fixed command bar unit tests

* Fixing spelling

* Adds unit tests for vcore quickstart feature code

* Fix lint
---
 .../NotebookTerminalComponent.test.tsx        |  49 ++++++++
 .../NotebookTerminalComponent.test.tsx.snap   |  24 ++++
 .../CommandBarComponentButtonFactory.test.ts  |  25 +++++
 .../Tabs/Shared/CheckFirewallRules.test.ts    |  91 +++++++++++++++
 src/Utils/EndpointValidation.ts               |  11 ++
 src/Utils/JestUtils.ts                        |   4 +
 src/Utils/NetworkUtility.test.ts              | 106 ++++++++++++++++++
 src/Utils/NetworkUtility.ts                   |  19 ++--
 src/userContext.test.ts                       |  18 +++
 9 files changed, 335 insertions(+), 12 deletions(-)
 create mode 100644 src/Explorer/Tabs/Shared/CheckFirewallRules.test.ts
 create mode 100644 src/Utils/JestUtils.ts
 create mode 100644 src/Utils/NetworkUtility.test.ts

diff --git a/src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx b/src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx
index fdc93adea..bd816ac31 100644
--- a/src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx
+++ b/src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx
@@ -32,6 +32,20 @@ const testCassandraAccount: DataModels.DatabaseAccount = {
   },
 };
 
+const testPostgresAccount: DataModels.DatabaseAccount = {
+  ...testAccount,
+  properties: {
+    postgresqlEndpoint: "https://testPostgresEndpoint.azure.com/",
+  },
+};
+
+const testVCoreMongoAccount: DataModels.DatabaseAccount = {
+  ...testAccount,
+  properties: {
+    vcoreMongoEndpoint: "https://testVCoreMongoEndpoint.azure.com/",
+  },
+};
+
 const testNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
   authToken: "authToken",
   notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com",
@@ -50,6 +64,18 @@ const testCassandraNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInf
   forwardingId: "Id",
 };
 
+const testPostgresNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
+  authToken: "authToken",
+  notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com/postgresql",
+  forwardingId: "Id",
+};
+
+const testVCoreMongoNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
+  authToken: "authToken",
+  notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com/mongovcore",
+  forwardingId: "Id",
+};
+
 describe("NotebookTerminalComponent", () => {
   it("renders terminal", () => {
     const props: NotebookTerminalComponentProps = {
@@ -94,4 +120,27 @@ describe("NotebookTerminalComponent", () => {
     const wrapper = shallow(<NotebookTerminalComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
+
+  it("renders Postgres shell", () => {
+    const props: NotebookTerminalComponentProps = {
+      databaseAccount: testPostgresAccount,
+      notebookServerInfo: testPostgresNotebookServerInfo,
+      tabId: undefined,
+    };
+
+    const wrapper = shallow(<NotebookTerminalComponent {...props} />);
+    expect(wrapper).toMatchSnapshot();
+  });
+
+  it("renders vCore Mongo shell", () => {
+    const props: NotebookTerminalComponentProps = {
+      databaseAccount: testVCoreMongoAccount,
+      notebookServerInfo: testVCoreMongoNotebookServerInfo,
+      tabId: undefined,
+      username: "username",
+    };
+
+    const wrapper = shallow(<NotebookTerminalComponent {...props} />);
+    expect(wrapper).toMatchSnapshot();
+  });
 });
diff --git a/src/Explorer/Controls/Notebook/__snapshots__/NotebookTerminalComponent.test.tsx.snap b/src/Explorer/Controls/Notebook/__snapshots__/NotebookTerminalComponent.test.tsx.snap
index a87a1c96d..440579141 100644
--- a/src/Explorer/Controls/Notebook/__snapshots__/NotebookTerminalComponent.test.tsx.snap
+++ b/src/Explorer/Controls/Notebook/__snapshots__/NotebookTerminalComponent.test.tsx.snap
@@ -1,5 +1,17 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`NotebookTerminalComponent renders Postgres shell 1`] = `
+<div
+  className="notebookTerminalContainer"
+>
+  <iframe
+    onLoad={[Function]}
+    src="terminal.html"
+    title="Terminal to Notebook Server"
+  />
+</div>
+`;
+
 exports[`NotebookTerminalComponent renders cassandra shell 1`] = `
 <div
   className="notebookTerminalContainer"
@@ -47,3 +59,15 @@ exports[`NotebookTerminalComponent renders terminal 1`] = `
   />
 </div>
 `;
+
+exports[`NotebookTerminalComponent renders vCore Mongo shell 1`] = `
+<div
+  className="notebookTerminalContainer"
+>
+  <iframe
+    onLoad={[Function]}
+    src="terminal.html"
+    title="Terminal to Notebook Server"
+  />
+</div>
+`;
diff --git a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.test.ts b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.test.ts
index 4e772ce30..cdf568d1a 100644
--- a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.test.ts
+++ b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.test.ts
@@ -343,6 +343,31 @@ describe("CommandBarComponentButtonFactory tests", () => {
     });
   });
 
+  describe("Open Postgres and vCore Mongo buttons", () => {
+    const openPostgresShellButtonLabel = "Open PSQL shell";
+    const openVCoreMongoShellButtonLabel = "Open MongoDB (vcore) shell";
+
+    beforeAll(() => {
+      mockExplorer = {} as Explorer;
+    });
+
+    it("creates Postgres shell button", () => {
+      const buttons = CommandBarComponentButtonFactory.createPostgreButtons(mockExplorer);
+      const openPostgresShellButton = buttons.find(
+        (button) => button.commandButtonLabel === openPostgresShellButtonLabel
+      );
+      expect(openPostgresShellButton).toBeDefined();
+    });
+
+    it("creates vCore Mongo shell button", () => {
+      const buttons = CommandBarComponentButtonFactory.createVCoreMongoButtons(mockExplorer);
+      const openVCoreMongoShellButton = buttons.find(
+        (button) => button.commandButtonLabel === openVCoreMongoShellButtonLabel
+      );
+      expect(openVCoreMongoShellButton).toBeDefined();
+    });
+  });
+
   describe("GitHub buttons", () => {
     const connectToGitHubBtnLabel = "Connect to GitHub";
     const manageGitHubSettingsBtnLabel = "Manage GitHub settings";
diff --git a/src/Explorer/Tabs/Shared/CheckFirewallRules.test.ts b/src/Explorer/Tabs/Shared/CheckFirewallRules.test.ts
new file mode 100644
index 000000000..e733d9e45
--- /dev/null
+++ b/src/Explorer/Tabs/Shared/CheckFirewallRules.test.ts
@@ -0,0 +1,91 @@
+import { DatabaseAccount, FirewallRule } from "Contracts/DataModels";
+import { checkFirewallRules } from "Explorer/Tabs/Shared/CheckFirewallRules";
+import { updateUserContext } from "UserContext";
+import { mockFunction } from "Utils/JestUtils";
+import { armRequest } from "Utils/arm/request";
+import React from "react";
+
+jest.mock("Utils/arm/request");
+const armRequestMock = mockFunction(armRequest);
+
+describe("CheckFirewallRule tests", () => {
+  const apiVersion = "2023-03-15-preview";
+  const rulePredicate = (rule: FirewallRule) =>
+    rule.properties.startIpAddress === "0.0.0.0" && rule.properties.endIpAddress === "255.255.255.255";
+  let isAllPublicIPAddressEnabled: boolean;
+  const setIsAllPublicIPAddressEnabled = jest.fn((value) => (isAllPublicIPAddressEnabled = value));
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const useStateMock: any = (initState: any) => [initState, setIsAllPublicIPAddressEnabled];
+  jest.spyOn(React, "useState").mockImplementation(useStateMock);
+
+  beforeAll(() => {
+    updateUserContext({
+      databaseAccount: {
+        id: "testResourceId",
+      } as DatabaseAccount,
+    });
+  });
+
+  it("returns 'all public IP addresses' is enabled for account with the proper firewall rule", async () => {
+    armRequestMock.mockResolvedValueOnce({
+      value: [
+        {
+          id: "resourceId",
+          name: "AllowAll",
+          type: "Microsoft.DocumentDB/mongoClusters/firewallRules",
+          properties: {
+            provisioningState: "Succeeded",
+            startIpAddress: "0.0.0.0",
+            endIpAddress: "255.255.255.255",
+          },
+        },
+      ],
+    });
+
+    await checkFirewallRules(apiVersion, rulePredicate, setIsAllPublicIPAddressEnabled);
+
+    expect(isAllPublicIPAddressEnabled).toBe(true);
+  });
+
+  it("returns 'all public IP addresses' is NOT enabled for account without the proper firewall rule", async () => {
+    armRequestMock.mockResolvedValueOnce([
+      {
+        id: "resourceId",
+        name: "AllowAll",
+        type: "Microsoft.DocumentDB/mongoClusters/firewallRules",
+        properties: {
+          provisioningState: "Succeeded",
+          startIpAddress: "10.10.10.10",
+          endIpAddress: "10.10.10.10",
+        },
+      },
+    ]);
+
+    await checkFirewallRules(apiVersion, rulePredicate, setIsAllPublicIPAddressEnabled);
+
+    expect(isAllPublicIPAddressEnabled).toBe(false);
+  });
+
+  it("sets message for account without the proper firewall rule", async () => {
+    armRequestMock.mockResolvedValueOnce([
+      {
+        id: "resourceId",
+        name: "AllowAll",
+        type: "Microsoft.DocumentDB/mongoClusters/firewallRules",
+        properties: {
+          provisioningState: "Succeeded",
+          startIpAddress: "0.0.0.0",
+          endIpAddress: "255.255.255.255",
+        },
+      },
+    ]);
+
+    const warningMessage = "This is a warning message";
+    let warningMessageResult: string;
+    const warningMessageFunc = (msg: string) => (warningMessageResult = msg);
+
+    await checkFirewallRules(apiVersion, rulePredicate, undefined, warningMessageFunc, warningMessage);
+
+    expect(warningMessageResult).toEqual(warningMessage);
+  });
+});
diff --git a/src/Utils/EndpointValidation.ts b/src/Utils/EndpointValidation.ts
index 1e7740367..d39c5892d 100644
--- a/src/Utils/EndpointValidation.ts
+++ b/src/Utils/EndpointValidation.ts
@@ -55,6 +55,17 @@ export const defaultAllowedBackendEndpoints: ReadonlyArray<string> = [
   "https://localhost:1234",
 ];
 
+export const PortalBackendIPs: { [key: string]: string[] } = {
+  "https://main.documentdb.ext.azure.com": ["104.42.195.92", "40.76.54.131"],
+  // DE doesn't talk to prod2 (main2) but it might be added
+  //"https://main2.documentdb.ext.azure.com": ["104.42.196.69"],
+  "https://main.documentdb.ext.azure.cn": ["139.217.8.252"],
+  "https://main.documentdb.ext.azure.us": ["52.244.48.71"],
+  // Add ussec and usnat when endpoint address is known:
+  //ussec: ["29.26.26.67", "29.26.26.66"],
+  //usnat: ["7.28.202.68"],
+};
+
 export const allowedMongoProxyEndpoints: ReadonlyArray<string> = [
   "https://main.documentdb.ext.azure.com",
   "https://main.documentdb.ext.azure.cn",
diff --git a/src/Utils/JestUtils.ts b/src/Utils/JestUtils.ts
new file mode 100644
index 000000000..899338e35
--- /dev/null
+++ b/src/Utils/JestUtils.ts
@@ -0,0 +1,4 @@
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function mockFunction<T extends (...args: any[]) => any>(fn: T): jest.MockedFunction<T> {
+  return fn as jest.MockedFunction<T>;
+}
diff --git a/src/Utils/NetworkUtility.test.ts b/src/Utils/NetworkUtility.test.ts
new file mode 100644
index 000000000..d33f13381
--- /dev/null
+++ b/src/Utils/NetworkUtility.test.ts
@@ -0,0 +1,106 @@
+import { resetConfigContext, updateConfigContext } from "ConfigContext";
+import { DatabaseAccount, IpRule } from "Contracts/DataModels";
+import { updateUserContext } from "UserContext";
+import { PortalBackendIPs } from "Utils/EndpointValidation";
+import { getNetworkSettingsWarningMessage } from "./NetworkUtility";
+
+describe("NetworkUtility tests", () => {
+  describe("getNetworkSettingsWarningMessage", () => {
+    const publicAccessMessagePart = "Please enable public access to proceed";
+    const accessMessagePart = "Please allow access from Azure Portal to proceed";
+    // validEnpoints are a subset of those from Utils/EndpointValidation/PortalBackendIPs
+    const validEndpoints = [
+      "https://main.documentdb.ext.azure.com",
+      "https://main.documentdb.ext.azure.cn",
+      "https://main.documentdb.ext.azure.us",
+    ];
+
+    let warningMessageResult: string;
+    const warningMessageFunc = (msg: string) => (warningMessageResult = msg);
+
+    beforeEach(() => {
+      warningMessageResult = undefined;
+    });
+
+    afterEach(() => {
+      resetConfigContext();
+    });
+
+    it("should return no message when publicNetworkAccess is enabled", async () => {
+      updateUserContext({
+        databaseAccount: {
+          properties: {
+            publicNetworkAccess: "Enabled",
+          },
+        } as DatabaseAccount,
+      });
+
+      await getNetworkSettingsWarningMessage(warningMessageFunc);
+      expect(warningMessageResult).toBeUndefined();
+    });
+
+    it("should return publicAccessMessage when publicNetworkAccess is disabled", async () => {
+      updateUserContext({
+        databaseAccount: {
+          properties: {
+            publicNetworkAccess: "Disabled",
+          },
+        } as DatabaseAccount,
+      });
+
+      await getNetworkSettingsWarningMessage(warningMessageFunc);
+      expect(warningMessageResult).toContain(publicAccessMessagePart);
+    });
+
+    it(`should return no message when the appropriate ip rules are added to mongo/cassandra account per endpoint`, () => {
+      validEndpoints.forEach(async (endpoint) => {
+        updateUserContext({
+          databaseAccount: {
+            kind: "MongoDB",
+            properties: {
+              ipRules: PortalBackendIPs[endpoint].map((ip: string) => ({ ipAddressOrRange: ip } as IpRule)),
+              publicNetworkAccess: "Enabled",
+            },
+          } as DatabaseAccount,
+        });
+
+        updateConfigContext({
+          BACKEND_ENDPOINT: endpoint,
+        });
+
+        let asyncWarningMessageResult: string;
+        const asyncWarningMessageFunc = (msg: string) => (asyncWarningMessageResult = msg);
+
+        await getNetworkSettingsWarningMessage(asyncWarningMessageFunc);
+        expect(asyncWarningMessageResult).toBeUndefined();
+      });
+    });
+
+    it("should return accessMessage when incorrent ip rule is added to mongo/cassandra account per endpoint", () => {
+      validEndpoints.forEach(async (endpoint) => {
+        updateUserContext({
+          databaseAccount: {
+            kind: "MongoDB",
+            properties: {
+              ipRules: [{ ipAddressOrRange: "1.1.1.1" }],
+              publicNetworkAccess: "Enabled",
+            },
+          } as DatabaseAccount,
+        });
+
+        updateConfigContext({
+          BACKEND_ENDPOINT: endpoint,
+        });
+
+        let asyncWarningMessageResult: string;
+        const asyncWarningMessageFunc = (msg: string) => (asyncWarningMessageResult = msg);
+
+        await getNetworkSettingsWarningMessage(asyncWarningMessageFunc);
+        expect(asyncWarningMessageResult).toContain(accessMessagePart);
+      });
+    });
+
+    // Postgres and vcore mongo account checks basically pass through to CheckFirewallRules so those
+    // tests are omitted here and included in CheckFirewallRules.test.ts
+  });
+});
diff --git a/src/Utils/NetworkUtility.ts b/src/Utils/NetworkUtility.ts
index 0037fec30..bda437dfb 100644
--- a/src/Utils/NetworkUtility.ts
+++ b/src/Utils/NetworkUtility.ts
@@ -1,15 +1,7 @@
+import { configContext } from "ConfigContext";
 import { checkFirewallRules } from "Explorer/Tabs/Shared/CheckFirewallRules";
 import { userContext } from "UserContext";
-
-const PortalIPs: { [key: string]: string[] } = {
-  prod1: ["104.42.195.92", "40.76.54.131"],
-  prod2: ["104.42.196.69"],
-  mooncake: ["139.217.8.252"],
-  blackforest: ["51.4.229.218"],
-  fairfax: ["52.244.48.71"],
-  ussec: ["29.26.26.67", "29.26.26.66"],
-  usnat: ["7.28.202.68"],
-};
+import { PortalBackendIPs } from "Utils/EndpointValidation";
 
 export const getNetworkSettingsWarningMessage = async (
   setStateFunc: (warningMessage: string) => void
@@ -28,6 +20,7 @@ export const getNetworkSettingsWarningMessage = async (
       setStateFunc,
       accessMessage
     );
+    return;
   } else if (userContext.apiType === "VCoreMongo") {
     checkFirewallRules(
       "2023-03-01-preview",
@@ -38,6 +31,7 @@ export const getNetworkSettingsWarningMessage = async (
       setStateFunc,
       accessMessage
     );
+    return;
   } else if (accountProperties) {
     // public network access is disabled
     if (
@@ -45,13 +39,14 @@ export const getNetworkSettingsWarningMessage = async (
       accountProperties.publicNetworkAccess !== "SecuredByPerimeter"
     ) {
       setStateFunc(publicAccessMessage);
+      return;
     }
 
     const ipRules = accountProperties.ipRules;
     // public network access is NOT set to "All networks"
-    if (ipRules.length > 0) {
+    if (ipRules?.length > 0) {
       if (userContext.apiType === "Cassandra" || userContext.apiType === "Mongo") {
-        const portalIPs = PortalIPs[userContext.portalEnv];
+        const portalIPs = PortalBackendIPs[configContext.BACKEND_ENDPOINT];
         let numberOfMatches = 0;
         ipRules.forEach((ipRule) => {
           if (portalIPs.indexOf(ipRule.ipAddressOrRange) !== -1) {
diff --git a/src/userContext.test.ts b/src/userContext.test.ts
index 0539633c1..672e80f5c 100644
--- a/src/userContext.test.ts
+++ b/src/userContext.test.ts
@@ -57,4 +57,22 @@ describe("shouldShowQueryPageOptions()", () => {
     });
     expect(userContext.apiType).toBe("Mongo");
   });
+
+  it("should be 'Postgres' for Postgres API", () => {
+    updateUserContext({
+      databaseAccount: {
+        kind: "Postgres",
+      } as DatabaseAccount,
+    });
+    expect(userContext.apiType).toBe("Postgres");
+  });
+
+  it("should be 'VCoreMongo' for vCore Mongo", () => {
+    updateUserContext({
+      databaseAccount: {
+        kind: "VCoreMongo",
+      } as DatabaseAccount,
+    });
+    expect(userContext.apiType).toBe("VCoreMongo");
+  });
 });