From 73e18584104a6a6149e7adab065282564b2ecf8c Mon Sep 17 00:00:00 2001
From: lironsh <lironsh@wix.com>
Date: Mon, 6 Jan 2025 17:20:54 +0200
Subject: [PATCH] test: add suitable test suites.

suitable test suites.

test: suitable test suites.

test: add suitable test suites

- Add `Copilot`, `StepPerformer`, `index` and `PromptCreator` suites
- Add utils file
- Add snapshot file for `PromptCreator` test
---
 src/Copilot.test.ts                           |  59 +++-
 src/actions/StepPerformer.test.ts             |  35 +++
 src/integration tests/index.test.ts           |  31 +-
 src/test-utils/expendAPICatalogUtils.ts       |  39 +++
 src/utils/CacheHandler.test.ts                |   8 +-
 src/utils/CodeEvaluator.test.ts               |   4 +-
 src/utils/PromptCreator.test.ts               |  22 +-
 .../__snapshots__/PromptCreator.test.ts.snap  | 270 ++++++++++++++++++
 src/utils/extractCodeBlock.test.ts            |   2 +-
 9 files changed, 452 insertions(+), 18 deletions(-)
 create mode 100644 src/test-utils/expendAPICatalogUtils.ts

diff --git a/src/Copilot.test.ts b/src/Copilot.test.ts
index 66ec9e4..c2e7424 100644
--- a/src/Copilot.test.ts
+++ b/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 {
+    customActionsCategory,
+    actionsCategory2,
+    actionsCategory,
+    tapButtonContext
+} from "./test-utils/expendAPICatalogUtils";
 
 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,43 @@ 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(actionsCategory);
+
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual(actionsCategory);
+            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(actionsCategory, tapButtonContext);
+
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual(actionsCategory);
+            expect(spyStepPerformer).toHaveBeenCalledWith(tapButtonContext);
+        });
+
+        it('should expend the API catalog with an existing category and an additional new category', () => {
+            Copilot.init(mockConfig);
+            const instance = Copilot.getInstance();
+
+            // item is added to existing category
+            instance.extendAPICatalog(actionsCategory)
+            instance.extendAPICatalog(actionsCategory2, tapButtonContext);
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual([...actionsCategory]);
+            expect(mockConfig.frameworkDriver.apiCatalog.categories[0].items).toHaveLength(2);
+            expect(spyStepPerformer).toHaveBeenCalledWith(tapButtonContext);
+
+            // another category is added
+            instance.extendAPICatalog(customActionsCategory);
+            expect(mockConfig.frameworkDriver.apiCatalog.categories).toEqual([...actionsCategory, ...customActionsCategory]);
+        });
+    });
 });
diff --git a/src/actions/StepPerformer.test.ts b/src/actions/StepPerformer.test.ts
index b9288a9..10664d0 100644
--- a/src/actions/StepPerformer.test.ts
+++ b/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 {functionContext, tapButtonContext, newFunctionContext} from "../test-utils/expendAPICatalogUtils";
 
 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(functionContext);
+
+            setupMocks();
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, functionContext);
+
+            // Extended context
+            const newFixedContext = { ...functionContext, ...tapButtonContext };
+            stepPerformer.extendJSContext(tapButtonContext);
+
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, newFixedContext);
+        });
+
+        it('should log when a context is overridden', async () => {
+            jest.spyOn(console, 'log');
+            stepPerformer.extendJSContext(functionContext);
+
+            setupMocks();
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, functionContext);
+
+            stepPerformer.extendJSContext(newFunctionContext);
+            expect(console.log).toHaveBeenCalledWith('Notice: Context function is overridden by the new context value');
+
+            await stepPerformer.perform(INTENT);
+            expect(mockCodeEvaluator.evaluate).toHaveBeenCalledWith(PROMPT_RESULT, newFunctionContext);
+
+        });
+    });
 });
diff --git a/src/integration tests/index.test.ts b/src/integration tests/index.test.ts
index da5a093..8040675 100644
--- a/src/integration tests/index.test.ts	
+++ b/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 {customActionsCategory, actionsCategory, tapButtonContext} from "../test-utils/expendAPICatalogUtils";
 
 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,27 @@ 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(customActionsCategory);
+            expect(spyPromptCreator).toHaveBeenCalledTimes(1);
+
+            copilot.extendAPICatalog(actionsCategory, tapButtonContext);
+            expect(spyPromptCreator).toHaveBeenCalledTimes(2);
+            expect(spyStepPerformer).toHaveBeenCalledTimes(1);
+        });
+    });
 });
diff --git a/src/test-utils/expendAPICatalogUtils.ts b/src/test-utils/expendAPICatalogUtils.ts
new file mode 100644
index 0000000..211cd8d
--- /dev/null
+++ b/src/test-utils/expendAPICatalogUtils.ts
@@ -0,0 +1,39 @@
+export const customActionsCategory = [{
+    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 actionsCategory = [{
+    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 actionsCategory2 = [{
+    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 tapButtonContext = {tapButton: 'new tap button'};
+export const functionContext = {function: 'newFlow'};
+export const newFunctionContext = {function: 'newFunction'};
diff --git a/src/utils/CacheHandler.test.ts b/src/utils/CacheHandler.test.ts
index 1b1f4d2..e17b16d 100644
--- a/src/utils/CacheHandler.test.ts
+++ b/src/utils/CacheHandler.test.ts
@@ -1,5 +1,5 @@
-import { CacheHandler } from './CacheHandler';
-import { mockCache, mockedCacheFile} from "../test-utils/cache";
+import {CacheHandler} from './CacheHandler';
+import {mockCache, mockedCacheFile} from "../test-utils/cache";
 
 jest.mock('fs');
 
@@ -13,7 +13,7 @@ describe('CacheHandler', () => {
 
     describe('cache and file operations', () => {
         it('should load cache from file successfully if the file exists and is valid', () => {
-            mockCache({ 'cacheKey': 'value' });
+            mockCache({'cacheKey': 'value'});
 
             expect(cacheHandler.getStepFromCache('cacheKey')).toBeUndefined();
 
@@ -28,7 +28,7 @@ describe('CacheHandler', () => {
             cacheHandler.addToTemporaryCache('cacheKey', 'value');
             cacheHandler.flushTemporaryCache();
 
-            expect(mockedCacheFile).toEqual({ 'cacheKey': 'value' });
+            expect(mockedCacheFile).toEqual({'cacheKey': 'value'});
         });
     });
 
diff --git a/src/utils/CodeEvaluator.test.ts b/src/utils/CodeEvaluator.test.ts
index 7652bbf..7fc9b2e 100644
--- a/src/utils/CodeEvaluator.test.ts
+++ b/src/utils/CodeEvaluator.test.ts
@@ -1,4 +1,4 @@
-import { CodeEvaluator } from '@/utils/CodeEvaluator';
+import {CodeEvaluator} from '@/utils/CodeEvaluator';
 
 describe('CodeEvaluator', () => {
     let codeEvaluator: CodeEvaluator;
@@ -26,7 +26,7 @@ describe('CodeEvaluator', () => {
         const contextVariable = 43;
         const validCode = 'return contextVariable - 1;';
 
-        await expect(codeEvaluator.evaluate(validCode, { contextVariable })).resolves.toStrictEqual({
+        await expect(codeEvaluator.evaluate(validCode, {contextVariable})).resolves.toStrictEqual({
             code: 'return contextVariable - 1;',
             result: 42
         });
diff --git a/src/utils/PromptCreator.test.ts b/src/utils/PromptCreator.test.ts
index f1b414f..3ded983 100644
--- a/src/utils/PromptCreator.test.ts
+++ b/src/utils/PromptCreator.test.ts
@@ -1,8 +1,9 @@
-import { PromptCreator } from './PromptCreator';
+import {PromptCreator} from './PromptCreator';
 import {
     PreviousStep,
     TestingFrameworkAPICatalog
 } from "@/types";
+import {customActionsCategory, actionsCategory2} from "../test-utils/expendAPICatalogUtils";
 
 const mockAPI: TestingFrameworkAPICatalog = {
     context: {},
@@ -93,4 +94,23 @@ describe('PromptCreator', () => {
 
         expect(prompt).toMatchSnapshot();
     });
+
+    describe('extentAPICategories', () => {
+        it('should extend the API catalog with new or exist category', () => {
+            // Add new category
+            const intent = 'expect button to be visible';
+            const viewHierarchy = '<View><Button testID="submit" title="Submit" /></View>';
+
+            promptCreator.extendAPICategories(customActionsCategory);
+            const promptNewCategory = promptCreator.createPrompt(intent, viewHierarchy, false, []);
+
+            expect(promptNewCategory).toMatchSnapshot();
+
+            // Add items to existing category
+            promptCreator.extendAPICategories(actionsCategory2);
+            const promptOldCategory = promptCreator.createPrompt(intent, viewHierarchy, false, []);
+
+            expect(promptOldCategory).toMatchSnapshot();
+        });
+    });
 });
diff --git a/src/utils/__snapshots__/PromptCreator.test.ts.snap b/src/utils/__snapshots__/PromptCreator.test.ts.snap
index 143f676..2a5c633 100644
--- a/src/utils/__snapshots__/PromptCreator.test.ts.snap
+++ b/src/utils/__snapshots__/PromptCreator.test.ts.snap
@@ -1,5 +1,275 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`PromptCreator extentAPICategories should extend the API catalog with new or exist category 1`] = `
+"# Test Code Generation
+
+You are an AI assistant tasked with generating test code for an application using the provided UI testing framework API.
+Please generate the minimal executable code to perform the desired intent based on the given information and context.
+
+## Context
+
+### Intent to perform
+
+Generate the minimal executable code to perform the following intent: "expect button to be visible"
+
+### View hierarchy
+
+\`\`\`
+<View><Button testID="submit" title="Submit" /></View>
+\`\`\`
+
+### Snapshot image
+
+No snapshot image is attached for this intent.
+
+## Available Testing Framework API
+
+### Actions
+
+#### tap(element: Element)
+
+Taps on the specified element.
+
+##### Example
+
+\`\`\`
+await element(by.id("button")).tap();
+\`\`\`
+
+##### Guidelines
+
+- Ensure the element is tappable before using this method.
+
+#### typeText(element: Element, text: string)
+
+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.
+
+### Assertions
+
+#### toBeVisible()
+
+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.
+
+### Matchers
+
+#### by.id(id: string)
+
+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.
+
+### Custom Actions
+
+#### swipe(direction: string)
+
+Swipes in the specified direction.
+
+##### Example
+
+\`\`\`
+await swipe("up");
+\`\`\`
+
+##### Guidelines
+
+- Use this method to scroll the screen.
+
+## Instructions
+
+Your task is to generate the minimal executable code to perform the following intent: "expect button to be visible"
+
+Please follow these steps carefully:
+
+1. Analyze the provided intent and the view hierarchy to understand the required action.
+2. Generate the minimal executable code required to perform the intent using the available API.
+3. If you cannot generate the relevant code due to ambiguity or invalid intent, return code that throws an informative error explaining the problem in one sentence.
+4. Wrap the generated code with backticks, without any additional formatting.
+5. Do not provide any additional code beyond the minimal executable code required to perform the intent.
+
+### Verify the prompt
+
+Before generating the code, please review the provided context and instructions to ensure they are clear and unambiguous. If you encounter any issues or have questions, please throw an informative error explaining the problem.
+
+### Examples
+
+#### Example of throwing an informative error:
+\`\`\`typescript
+throw new Error("Unable to find the 'Submit' button element in the current context.");
+\`\`\`
+
+Please provide your response below:"
+`;
+
+exports[`PromptCreator extentAPICategories should extend the API catalog with new or exist category 2`] = `
+"# Test Code Generation
+
+You are an AI assistant tasked with generating test code for an application using the provided UI testing framework API.
+Please generate the minimal executable code to perform the desired intent based on the given information and context.
+
+## Context
+
+### Intent to perform
+
+Generate the minimal executable code to perform the following intent: "expect button to be visible"
+
+### View hierarchy
+
+\`\`\`
+<View><Button testID="submit" title="Submit" /></View>
+\`\`\`
+
+### Snapshot image
+
+No snapshot image is attached for this intent.
+
+## Available Testing Framework API
+
+### Actions
+
+#### tap(element: Element)
+
+Taps on the specified element.
+
+##### Example
+
+\`\`\`
+await element(by.id("button")).tap();
+\`\`\`
+
+##### Guidelines
+
+- Ensure the element is tappable before using this method.
+
+#### typeText(element: Element, text: string)
+
+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.
+
+#### swipe(direction: string)
+
+Swipes in the specified direction.
+
+##### Example
+
+\`\`\`
+await swipe("up");
+\`\`\`
+
+##### Guidelines
+
+- Use this method to scroll the screen.
+
+### Assertions
+
+#### toBeVisible()
+
+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.
+
+### Matchers
+
+#### by.id(id: string)
+
+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.
+
+### Custom Actions
+
+#### swipe(direction: string)
+
+Swipes in the specified direction.
+
+##### Example
+
+\`\`\`
+await swipe("up");
+\`\`\`
+
+##### Guidelines
+
+- Use this method to scroll the screen.
+
+## Instructions
+
+Your task is to generate the minimal executable code to perform the following intent: "expect button to be visible"
+
+Please follow these steps carefully:
+
+1. Analyze the provided intent and the view hierarchy to understand the required action.
+2. Generate the minimal executable code required to perform the intent using the available API.
+3. If you cannot generate the relevant code due to ambiguity or invalid intent, return code that throws an informative error explaining the problem in one sentence.
+4. Wrap the generated code with backticks, without any additional formatting.
+5. Do not provide any additional code beyond the minimal executable code required to perform the intent.
+
+### Verify the prompt
+
+Before generating the code, please review the provided context and instructions to ensure they are clear and unambiguous. If you encounter any issues or have questions, please throw an informative error explaining the problem.
+
+### Examples
+
+#### Example of throwing an informative error:
+\`\`\`typescript
+throw new Error("Unable to find the 'Submit' button element in the current context.");
+\`\`\`
+
+Please provide your response below:"
+`;
+
 exports[`PromptCreator should create a prompt for an intent correctly 1`] = `
 "# Test Code Generation
 
diff --git a/src/utils/extractCodeBlock.test.ts b/src/utils/extractCodeBlock.test.ts
index 4762ae9..cd9f953 100644
--- a/src/utils/extractCodeBlock.test.ts
+++ b/src/utils/extractCodeBlock.test.ts
@@ -1,4 +1,4 @@
-import { extractCodeBlock } from "./extractCodeBlock";
+import {extractCodeBlock} from "./extractCodeBlock";
 
 describe('extractCodeBlock', () => {
     const runExtractionTest = (input: string, expected: string) => {