From af6975fbe576f49b9e72f990263c86e76ad801a6 Mon Sep 17 00:00:00 2001
From: Priyansh Garg <priyanshgarg30@gmail.com>
Date: Mon, 11 Nov 2024 23:57:05 +0530
Subject: [PATCH] Add types for custom commands on page object sections.
 (#4298)

---
 types/index.d.ts                  |  4 ++++
 types/page-object.d.ts            |  3 ++-
 types/tests/page-object.test-d.ts | 40 +++++++++++++++++++++++++++----
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/types/index.d.ts b/types/index.d.ts
index 2d7254ec5..b56a057ff 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -68,6 +68,10 @@ export type NightwatchGenericCallback<T> = (
 
 export type Awaitable<T, V> = Omit<T, 'then'> & PromiseLike<V>;
 
+export type KeysFilter<T, U> = {
+  [K in keyof T]-?: T[K] extends U ? K : never;
+}[keyof T];
+
 // tslint:disable-next-line
 type VoidToNull<T> = T extends void ? null : T;
 
diff --git a/types/page-object.d.ts b/types/page-object.d.ts
index b6420a53e..3e15e1556 100644
--- a/types/page-object.d.ts
+++ b/types/page-object.d.ts
@@ -9,8 +9,8 @@ import {
   ElementCommands,
   ElementFunction,
   Expect,
+  KeysFilter,
   LocateStrategy,
-  NamespacedApi,
   NightwatchAPI,
   NightwatchClient,
   NightwatchComponentTestingCommands,
@@ -147,6 +147,7 @@ export type EnhancedSectionInstance<
   Commands &
   ElementCommands &
   ChromiumClientCommands &
+  Pick<NightwatchCustomCommands, KeysFilter<NightwatchCustomCommands, Function>> & // eslint-disable-line @typescript-eslint/ban-types
   Pick<
     NightwatchComponentTestingCommands,
     'importScript' | 'launchComponentRenderer' | 'mountComponent'
diff --git a/types/tests/page-object.test-d.ts b/types/tests/page-object.test-d.ts
index 35becd308..17cfbffae 100644
--- a/types/tests/page-object.test-d.ts
+++ b/types/tests/page-object.test-d.ts
@@ -1,21 +1,33 @@
-// Page object file
-import {EnhancedPageObject, PageObjectModel} from '..';
+import {expectError, expectType} from 'tsd';
+import {EnhancedPageObject, PageObjectModel, Awaitable, NightwatchAPI} from '..';
 
+// Page object file
 const fileUploadPageElements = {
   fileUploadInput: 'input#file-upload',
   submitButton: 'input#file-submit',
   uploadFiles: '#uploaded-files'
 };
 
+const menuSection = {
+  selector: 'nav',
+  elements: {
+    home: 'a[href="/"]',
+    about: 'a[href="/about"]'
+  }
+};
+
 const fileUploadPage = {
   url(this: EnhancedPageObject) {
     return `${this.api.launch_url}/upload`;
   },
-  elements: fileUploadPageElements
+  elements: fileUploadPageElements,
+  sections: {
+    menu: menuSection
+  }
 } satisfies PageObjectModel;
 
 export interface FileUploadPage extends
-  EnhancedPageObject<{}, typeof fileUploadPageElements, {}, {}, () => string> {} // eslint-disable-line @typescript-eslint/ban-types
+  EnhancedPageObject<{}, typeof fileUploadPageElements, typeof fileUploadPage.sections, {}, () => string> {} // eslint-disable-line @typescript-eslint/ban-types
 
 export default fileUploadPage;
 
@@ -26,6 +38,13 @@ declare module '..' {
   interface NightwatchCustomPageObjects {
     FileUpload(): FileUploadPage;
   }
+  interface NightwatchCustomCommands {
+    uploadFile1(selector: string, filePath: string): NightwatchAPI;
+    uploadFile2(selector: string, filePath: string): Awaitable<NightwatchAPI, null>;
+    upload: {
+      file(selector: string, filePath: string): NightwatchAPI;
+    }
+  }
 }
 
 
@@ -42,6 +61,19 @@ describe('File Upload', function() {
       .click(fileUploadPage.elements.submitButton)
       .expect.element('@uploadFiles').text.to.equal('test.txt');
 
+    // test custom commands over page object
+    expectType<NightwatchAPI>(fileUploadPage.uploadFile1('@fileUploadInput', 'test2.txt'));
+    expectType<Awaitable<NightwatchAPI, null>>(fileUploadPage.uploadFile2('@fileUploadInput', 'test2.txt'));
+    // should fail ideally but succeeding
+    expectType<NightwatchAPI>(fileUploadPage.upload.file('@fileUploadInput', 'test2.txt'));
+
+    // test custom commands over page object sections
+    const menuSection = fileUploadPage.section.menu;
+    expectType<NightwatchAPI>(menuSection.uploadFile1('@fileUploadInput', 'test2.txt'));
+    expectType<Awaitable<NightwatchAPI, null>>(menuSection.uploadFile2('@fileUploadInput', 'test2.txt'));
+    // should fail because the namespaces from custom commands are not loaded directly into the section
+    expectError(menuSection.upload.file('@fileUploadInput', 'test2.txt'));
+
     browser.end();
   });
 });