Merge pull request #34 from wix-incubator/Introduce-extendAPICatalog-…

…method

Introduce extend api catalog method

src/Copilot.test.ts

@@ -1,9 +1,14 @@
-import { Copilot } from '@/Copilot';
-import { StepPerformer } from '@/actions/StepPerformer';
-import { CopilotError } from '@/errors/CopilotError';
-import { Config } from "@/types";
-import fs from "fs";
-import { mockCache, mockedCacheFile } from "./test-utils/cache";
+import {Copilot} from '@/Copilot';
+import {StepPerformer} from '@/actions/StepPerformer';
+import {CopilotError} from '@/errors/CopilotError';
+import {Config} from "@/types";
+import {mockCache, mockedCacheFile} from "./test-utils/cache";
+import {
+    bazCategory,
+    barCategory2,
+    barCategory1,
+    dummyContext
+} from "./test-utils/APICatalogTestUtils";
 
 jest.mock('@/actions/StepPerformer');
 jest.mock('fs');
@@ -28,7 +33,8 @@ describe('Copilot', () => {
                 isSnapshotImageSupported: jest.fn().mockReturnValue(true)
             }
         };
-        jest.spyOn(console, 'error').mockImplementation(() => {});
+        jest.spyOn(console, 'error').mockImplementation(() => {
+        });
 
         (StepPerformer.prototype.perform as jest.Mock).mockResolvedValue({code: 'code', result: true});
     });
@@ -188,4 +194,49 @@ describe('Copilot', () => {
             expect(mockedCacheFile).toBeUndefined();
         });
     });
+
+    describe('extend API catalog', () => {
+        const spyStepPerformer = jest.spyOn(StepPerformer.prototype, 'extendJSContext');
+        it('should extend the API catalog with a new category', () => {
+            Copilot.init(mockConfig);
+            const instance = Copilot.getInstance();
+
+            instance.extendAPICatalog([barCategory1]);
+
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual([barCategory1]);
+            expect(spyStepPerformer).not.toHaveBeenCalled();
+
+        });
+
+        it('should extend the API catalog with a new category and context', () => {
+            Copilot.init(mockConfig);
+            const instance = Copilot.getInstance();
+            instance.extendAPICatalog([barCategory1], dummyContext);
+
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual([barCategory1]);
+            expect(spyStepPerformer).toHaveBeenCalledWith(dummyContext);
+        });
+
+        it('should extend the API catalog with an existing category', () => {
+            Copilot.init(mockConfig);
+            const instance = Copilot.getInstance();
+
+            instance.extendAPICatalog([barCategory1])
+            instance.extendAPICatalog([barCategory2], dummyContext);
+
+            expect(mockConfig.frameworkDriver.apiCatalog.categories.length).toEqual(1);
+            expect(mockConfig.frameworkDriver.apiCatalog.categories[0].items).toEqual([...barCategory1.items, ...barCategory2.items]);
+            expect(spyStepPerformer).toHaveBeenCalledWith(dummyContext);
+        });
+
+        it('should extend the API catalog with a new category', () => {
+            Copilot.init(mockConfig);
+            const instance = Copilot.getInstance();
+
+            instance.extendAPICatalog([barCategory1]);
+            instance.extendAPICatalog([bazCategory]);
+
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual([barCategory1, bazCategory]);
+        });
+    });
 });

src/Copilot.ts

@@ -3,7 +3,7 @@ import {PromptCreator} from "@/utils/PromptCreator";
 import {CodeEvaluator} from "@/utils/CodeEvaluator";
 import {SnapshotManager} from "@/utils/SnapshotManager";
 import {StepPerformer} from "@/actions/StepPerformer";
-import {Config, PreviousStep} from "@/types";
+import {Config, PreviousStep, TestingFrameworkAPICatalogCategory} from "@/types";
 import {CacheHandler} from "@/utils/CacheHandler";
 
 /**
@@ -109,6 +109,17 @@ export class Copilot {
             this.cacheHandler.flushTemporaryCache();
     }
 
+    /**
+     * Enriches the API catalog by adding the provided categories and JS context.
+     * @param categories - The categories to register.
+     * @param context - (Optional) Additional JS context to register.
+     */
+    extendAPICatalog(categories: TestingFrameworkAPICatalogCategory[], context?: any): void {
+        this.promptCreator.extendAPICategories(categories);
+        if (context)
+            this.stepPerformer.extendJSContext(context);
+    }
+
     private didPerformStep(step: string, code: string, result: any): void {
         this.previousSteps = [...this.previousSteps, {
             step,

src/actions/StepPerformer.test.ts

@@ -5,6 +5,7 @@ import {SnapshotManager} from '@/utils/SnapshotManager';
 import {CacheHandler} from '@/utils/CacheHandler';
 import {PromptHandler, TestingFrameworkAPICatalog} from '@/types';
 import * as crypto from 'crypto';
+import {dummyContext, dummyBarContext1, dummyBarContext2} from "../test-utils/APICatalogTestUtils";
 
 jest.mock('fs');
 jest.mock('crypto');
@@ -283,4 +284,38 @@ describe('StepPerformer', () => {
         expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith('generated code', mockContext);
         expect(mockCacheHandler.addToTemporaryCache).toHaveBeenCalled();
     });
+
+    describe('extendJSContext', () =>{
+        it('should extend the context with the given object', async () => {
+            // Initial context
+            stepPerformer.extendJSContext(dummyBarContext1);
+
+            setupMocks();
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, dummyBarContext1);
+
+            // Extended context
+            const extendedContext = { ...dummyBarContext1, ...dummyContext };
+            stepPerformer.extendJSContext(dummyContext);
+
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, extendedContext);
+        });
+
+        it('should log when a context is overridden', async () => {
+            jest.spyOn(console, 'log');
+            stepPerformer.extendJSContext(dummyBarContext1);
+
+            setupMocks();
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, dummyBarContext1);
+
+            stepPerformer.extendJSContext(dummyBarContext2);
+            expect(console.log).toHaveBeenCalledWith('Notice: Context bar is overridden by the new context value');
+
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, dummyBarContext2);
+
+        });
+    });
 });

src/actions/StepPerformer.ts

@@ -19,9 +19,19 @@ export class StepPerformer {
     ) {
     }
 
+    extendJSContext(newContext: any): void {
+        for (const key in newContext) {
+            if (key in this.context) {
+                console.log(`Notice: Context ${key} is overridden by the new context value`);
+                break;
+            }
+        }
+        this.context = {...this.context, ...newContext};
+    }
+
     private generateCacheKey(step: string, previous: PreviousStep[], viewHierarchy: string): string {
         const viewHierarchyHash = crypto.createHash('md5').update(viewHierarchy).digest('hex');
-        return JSON.stringify({ step, previous, viewHierarchyHash });
+        return JSON.stringify({step, previous, viewHierarchyHash});
     }
 
     private async captureSnapshotAndViewHierarchy() {
@@ -32,7 +42,7 @@ export class StepPerformer {
 
         const isSnapshotImageAttached = snapshot != null && this.promptHandler.isSnapshotImageSupported();
 
-        return { snapshot, viewHierarchy, isSnapshotImageAttached };
+        return {snapshot, viewHierarchy, isSnapshotImageAttached};
     }
 
     private shouldOverrideCache() {
@@ -73,7 +83,7 @@ export class StepPerformer {
 
         for (let attempt = 1; attempt <= attempts; attempt++) {
             try {
-                const { snapshot, viewHierarchy, isSnapshotImageAttached } = await this.captureSnapshotAndViewHierarchy();
+                const {snapshot, viewHierarchy, isSnapshotImageAttached} = await this.captureSnapshotAndViewHierarchy();
 
                 const code = await this.generateCode(step, previous, snapshot, viewHierarchy, isSnapshotImageAttached);
                 lastCode = code;

src/index.ts

@@ -1,5 +1,5 @@
 import {Copilot} from "@/Copilot";
-import {CopilotFacade, Config} from "@/types";
+import {CopilotFacade, Config, TestingFrameworkAPICatalogCategory} from "@/types";
 
 const copilot: CopilotFacade = {
     init: (config: Config) => {
@@ -23,6 +23,9 @@ const copilot: CopilotFacade = {
         }
 
         return result;
+    },
+    extendAPICatalog: (categories: TestingFrameworkAPICatalogCategory[], context?: any) => {
+        Copilot.getInstance().extendAPICatalog(categories, context);
     }
 };
 

src/integration tests/index.test.ts

@@ -1,9 +1,12 @@
 import copilot from "@/index";
 import fs from 'fs';
-import { Copilot } from "@/Copilot";
-import { PromptHandler, TestingFrameworkDriver } from "@/types";
+import {Copilot} from "@/Copilot";
+import {PromptHandler, TestingFrameworkDriver} from "@/types";
 import * as crypto from 'crypto';
 import {mockedCacheFile, mockCache} from "../test-utils/cache";
+import {PromptCreator} from "../utils/PromptCreator";
+import {StepPerformer} from "../actions/StepPerformer";
+import {bazCategory, barCategory1, dummyContext} from "../test-utils/APICatalogTestUtils";
 
 jest.mock('crypto');
 jest.mock('fs');
@@ -79,7 +82,6 @@ describe('Copilot Integration Tests', () => {
 
         it('should successfully perform an action', async () => {
             mockPromptHandler.runPrompt.mockResolvedValue('// No operation');
-
             await expect(copilot.perform('Tap on the login button')).resolves.not.toThrow();
 
             expect(mockFrameworkDriver.captureSnapshotImage).toHaveBeenCalled();
@@ -330,4 +332,28 @@ describe('Copilot Integration Tests', () => {
             );
         });
     });
+
+    describe('API Catalog Extension', () => {
+        const spyPromptCreator = jest.spyOn(PromptCreator.prototype, 'extendAPICategories');
+        const spyStepPerformer = jest.spyOn(StepPerformer.prototype, 'extendJSContext');
+
+        beforeEach(() => {
+            jest.clearAllMocks();
+            copilot.init({
+                frameworkDriver: mockFrameworkDriver,
+                promptHandler: mockPromptHandler,
+            });
+            copilot.start();
+        });
+
+        it('should call relevant functions to extend the catalog', () => {
+
+            copilot.extendAPICatalog([bazCategory]);
+            expect(spyPromptCreator).toHaveBeenCalledTimes(1);
+
+            copilot.extendAPICatalog([barCategory1], dummyContext);
+            expect(spyPromptCreator).toHaveBeenCalledTimes(2);
+            expect(spyStepPerformer).toHaveBeenCalledTimes(1);
+        });
+    });
 });

src/test-utils/APICatalogTestUtils.ts

@@ -0,0 +1,102 @@
+import {TestingFrameworkAPICatalog} from "@/types";
+
+export const bazCategory = {
+    title: 'Custom Actions',
+    items: [
+        {
+            signature: 'swipe(direction: string)',
+            description: 'Swipes in the specified direction.',
+            example: 'await swipe("up");',
+            guidelines: ['Use this method to scroll the screen.']
+        }
+    ]
+};
+
+export const barCategory1 = {
+    title: 'Actions',
+    items: [
+        {
+            signature: 'tapButton(id: string)',
+            description: 'Taps the button with the specified ID.',
+            example: 'await tapButton("submit");',
+            guidelines: ['Use this method to tap buttons.']
+        }
+    ]
+};
+
+export const barCategory2 = {
+    title: 'Actions',
+    items: [
+        {
+            signature: 'swipe(direction: string)',
+            description: 'Swipes in the specified direction.',
+            example: 'await swipe("up");',
+            guidelines: ['Use this method to scroll the screen.']
+        }
+    ]
+};
+
+export const dummyContext = {foo: jest.fn()};
+export const dummyBarContext1 = {bar: jest.fn()};
+export const dummyBarContext2 = {bar: jest.fn()};
+
+export const promptCreatorConstructorMockAPI: TestingFrameworkAPICatalog = {
+    context: {},
+    categories: [
+        {
+            title: 'Actions',
+            items: [{
+                signature: 'tap(element: Element)',
+                description: 'Taps on the specified element.',
+                example: 'await element(by.id("button")).tap();',
+                guidelines: ['Ensure the element is tappable before using this method.']
+            },
+                {
+                    signature: 'typeText(element: Element, text: string)',
+                    description: 'Types the specified text into the element.',
+                    example: 'await element(by.id("input")).typeText("Hello, World!");',
+                    guidelines: ['Use this method only on text input elements.']
+                }]
+        },
+        {
+            title: 'Assertions',
+            items: [{
+                signature: 'toBeVisible()',
+                description: 'Asserts that the element is visible on the screen.',
+                example: 'await expect(element(by.id("title"))).toBeVisible();',
+                guidelines: ['Consider scroll position when using this assertion.']
+            }]
+        },
+        {
+            title: 'Assertions',
+            items: [{
+                signature: 'toBeEnabled()',
+                description: 'Asserts that the element is enabled and can be interacted with.',
+                example: 'await expect(element(by.id("submitButton"))).toBeEnabled();',
+                guidelines: ['Ensure that the element is not disabled before performing actions.']
+            }]
+        },
+        {
+            title: 'Matchers',
+            items: [
+                {
+                    signature: 'by.id(id: string)',
+                    description: 'Matches elements by their ID attribute.',
+                    example: 'element(by.id("uniqueId"))',
+                    guidelines: ['Use unique IDs for elements to avoid conflicts, combine with atIndex() if necessary.']
+                }
+            ]
+        },
+        {
+            title: 'Actions',
+            items: [
+                {
+                    signature: 'swipe(direction: string)',
+                    description: 'Swipes in the specified direction.',
+                    example: 'await swipe("up");',
+                    guidelines: ['Use this method to scroll the screen.']
+                }
+            ]
+        }
+    ]
+};

src/types.ts

@@ -41,6 +41,14 @@ export interface CopilotFacade {
      * @example 'Tap on the login button', 'A login form should be visible'
      */
     perform: (...steps: string[]) => Promise<string>;
+
+    /**
+     * Extends the API catalog of the testing framework with additional APIs (categories and JS context).
+     * @param context The variables of the testing framework (i.e. exposes the matching function, expect, etc.).
+     * @param categories The categories to add to the API catalog.
+     * @note This can be used to add custom categories and items to the API catalog.
+     */
+    extendAPICatalog: (categories: TestingFrameworkAPICatalogCategory[], context?: any,) => void;
 }
 
 /**