From 68c42fb36150834c70c39bf47b964a2b58c57c1a Mon Sep 17 00:00:00 2001
From: Sung-Hyun Kang <sunghyunkang@microsoft.com>
Date: Thu, 24 Apr 2025 23:40:05 -0500
Subject: [PATCH] Show system partition key value and add test cases

---
 .github/workflows/ci.yml                      |  4 ++--
 .../Controls/InputDataList/InputDataList.tsx  |  1 +
 .../Tabs/DocumentsTabV2/DocumentsTabV2.tsx    |  9 +++++++--
 test/fx.ts                                    | 10 ++++++++++
 test/mongo/document.spec.ts                   |  4 ++++
 test/sql/document.spec.ts                     |  9 +++++----
 test/sql/testCases.ts                         | 19 ++++++++++++++++++-
 test/testData.ts                              |  1 +
 8 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 56ed46edc..54f784668 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -164,8 +164,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        shardIndex: [1, 2, 3, 4, 5, 6, 7, 8]
-        shardTotal: [8]
+        shardIndex: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
+        shardTotal: [12]
     steps:
       - uses: actions/checkout@v4
       - name: "Az CLI login"
diff --git a/src/Explorer/Controls/InputDataList/InputDataList.tsx b/src/Explorer/Controls/InputDataList/InputDataList.tsx
index cd31db53b..2f483d362 100644
--- a/src/Explorer/Controls/InputDataList/InputDataList.tsx
+++ b/src/Explorer/Controls/InputDataList/InputDataList.tsx
@@ -193,6 +193,7 @@ export const InputDataList: FC<InputDataListProps> = ({
     <>
       <Input
         id="filterInput"
+        data-test={"DocumentsTab/FilterInput"}
         ref={inputRef}
         type="text"
         size="small"
diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
index 4c96064c0..58b976d44 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
@@ -773,8 +773,11 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     [_collection, _partitionKey],
   );
   const partitionKeyPropertyHeaders: string[] = useMemo(
-    () => (partitionKey?.systemKey ? [] : _collection?.partitionKeyPropertyHeaders || partitionKey?.paths),
-    [_collection?.partitionKeyPropertyHeaders, partitionKey?.paths, partitionKey?.systemKey],
+    () =>
+      isPreferredApiMongoDB && partitionKey?.systemKey
+        ? []
+        : _collection?.partitionKeyPropertyHeaders || partitionKey?.paths,
+    [_collection?.partitionKeyPropertyHeaders, partitionKey?.paths, partitionKey?.systemKey, isPreferredApiMongoDB],
   );
   let partitionKeyProperties = useMemo(() => {
     return partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
@@ -2116,6 +2119,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
           />
           <Button
             appearance="primary"
+            data-test={"DocumentsTab/ApplyFilter"}
             size="small"
             onClick={() => {
               if (isExecuting) {
@@ -2188,6 +2192,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
               {tableItems.length > 0 && (
                 <a
                   className={styles.loadMore}
+                  data-test={"DocumentsTab/LoadMore"}
                   role="button"
                   tabIndex={0}
                   onClick={() => loadNextPage(documentsIterator.iterator, false)}
diff --git a/test/fx.ts b/test/fx.ts
index 30c73e397..e15513004 100644
--- a/test/fx.ts
+++ b/test/fx.ts
@@ -223,6 +223,9 @@ export class DocumentsTab {
   documentsListPane: Locator;
   documentResultsPane: Locator;
   resultsEditor: Editor;
+  loadMoreButton: Locator;
+  filterInput: Locator;
+  filterButton: Locator;
 
   constructor(
     public frame: Frame,
@@ -234,6 +237,13 @@ export class DocumentsTab {
     this.documentsListPane = this.locator.getByTestId("DocumentsTab/DocumentsPane");
     this.documentResultsPane = this.locator.getByTestId("DocumentsTab/ResultsPane");
     this.resultsEditor = new Editor(this.frame, this.documentResultsPane.getByTestId("EditorReact/Host/Loaded"));
+    this.loadMoreButton = this.documentsListPane.getByTestId("DocumentsTab/LoadMore");
+    this.filterInput = this.documentsFilter.getByTestId("DocumentsTab/FilterInput");
+    this.filterButton = this.documentsFilter.getByTestId("DocumentsTab/ApplyFilter");
+  }
+
+  async setFilter(text: string) {
+    await this.filterInput.fill(text);
   }
 }
 
diff --git a/test/mongo/document.spec.ts b/test/mongo/document.spec.ts
index 34d27ee21..747e63783 100644
--- a/test/mongo/document.spec.ts
+++ b/test/mongo/document.spec.ts
@@ -68,8 +68,12 @@ for (const { name, databaseId, containerId, documents } of documentTestCases) {
             await documentsTab.resultsEditor.setText(JSON.stringify(newDocument));
             const saveButton = await explorer.waitForCommandBarButton("Save", 5000);
             await saveButton.click({ timeout: 5000 });
+            await expect(saveButton).toBeHidden({ timeout: 5000 });
           }, 3);
 
+          await documentsTab.setFilter(`{_id: "${newDocumentId}"}`);
+          await documentsTab.filterButton.click();
+
           const newSpan = documentsTab.documentsListPane.getByText(newDocumentId, { exact: true }).nth(0);
           await newSpan.waitFor();
           await newSpan.click();
diff --git a/test/sql/document.spec.ts b/test/sql/document.spec.ts
index 095b47c6a..458580d3e 100644
--- a/test/sql/document.spec.ts
+++ b/test/sql/document.spec.ts
@@ -51,10 +51,6 @@ for (const { name, databaseId, containerId, documents } of documentTestCases) {
           let newDocumentId;
           await page.waitForTimeout(5000);
           await retry(async () => {
-            // const discardButton = await explorer.waitForCommandBarButton("Discard", 5000);
-            // if (await discardButton.isEnabled()) {
-            //   await discardButton.click();
-            // }
             const newDocumentButton = await explorer.waitForCommandBarButton("New Item", 5000);
             await expect(newDocumentButton).toBeVisible();
             await expect(newDocumentButton).toBeEnabled();
@@ -72,10 +68,15 @@ for (const { name, databaseId, containerId, documents } of documentTestCases) {
             await documentsTab.resultsEditor.setText(JSON.stringify(newDocument));
             const saveButton = await explorer.waitForCommandBarButton("Save", 5000);
             await saveButton.click({ timeout: 5000 });
+            await expect(saveButton).toBeHidden({ timeout: 5000 });
           }, 3);
 
+          await documentsTab.setFilter(`WHERE c.id = "${newDocumentId}"`);
+          await documentsTab.filterButton.click();
+
           const newSpan = documentsTab.documentsListPane.getByText(newDocumentId, { exact: true }).nth(0);
           await newSpan.waitFor();
+
           await newSpan.click();
           await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
 
diff --git a/test/sql/testCases.ts b/test/sql/testCases.ts
index 8c4c3178f..4398c93eb 100644
--- a/test/sql/testCases.ts
+++ b/test/sql/testCases.ts
@@ -5,7 +5,24 @@ export const documentTestCases: DocumentTestCase[] = [
     name: "System Partition Key",
     databaseId: "e2etests-sql-readonly",
     containerId: "systemPartitionKey",
-    documents: [{ documentId: "systempartition", partitionKeys: [] }],
+    documents: [
+      {
+        documentId: "systempartition",
+        partitionKeys: [{ key: "/_partitionKey", value: "partitionKey" }],
+        skipCreateDelete: true,
+      },
+      {
+        documentId: "systempartition_empty",
+        partitionKeys: [{ key: "/_partitionKey", value: "" }],
+        skipCreateDelete: true,
+      },
+      {
+        documentId: "systempartition_null",
+        partitionKeys: [{ key: "/_partitionKey", value: null }],
+        skipCreateDelete: true,
+      },
+      { documentId: "systempartition_missing", partitionKeys: [] },
+    ],
   },
   {
     name: "Single Partition Key",
diff --git a/test/testData.ts b/test/testData.ts
index 3af3c903c..faa7b0d25 100644
--- a/test/testData.ts
+++ b/test/testData.ts
@@ -26,6 +26,7 @@ export interface DocumentTestCase {
 export interface TestDocument {
   documentId: string;
   partitionKeys?: PartitionKey[];
+  skipCreateDelete?: boolean;
 }
 
 export interface PartitionKey {