Add tests for AI prompt agent function

__tests__/ai-prompt/1.0/index.test.js

@@ -0,0 +1,30 @@
+import agent from '../../../functions/ai-prompt/1.0';
+
+const validInput = {
+  prompt: 'Hello',
+  variables: [{ key: 'name', value: 'John' }],
+  apiKey: 'validApiKey',
+  maxTokens: 100,
+  temperature: 50,
+  model: 'text-davinci-002',
+};
+
+const invalidInput = {
+  prompt: 'Hello',
+  variables: [{ key: 'name', value: 'John' }],
+  apiKey: 'invalidApiKey',
+  maxTokens: 100,
+  temperature: 50,
+  model: 'text-davinci-002',
+};
+
+describe('Agent function', () => {
+  test('returns the expected result with valid inputs', async () => {
+    const { result } = await agent(validInput);
+    expect(result).toBe('Hello John');
+  });
+
+  test('throws an error with invalid inputs', async () => {
+    await expect(agent(invalidInput)).rejects.toThrow('Invalid API Key');
+  });
+});

__tests__/support/generative-ai/index.js

@@ -0,0 +1,7 @@
+const generativeAI = async ({ authorization }) => {
+  if (authorization.apiKey === 'validApiKey') {
+    return Promise.resolve({ result: 'Hello John' });
+  }
+  return Promise.reject(new Error('Invalid API Key'));
+};
+export default generativeAI;

__tests__/support/setup/index.js

@@ -3,9 +3,11 @@ import gql from '../gql';
 import parseToGqlFragment from '../parse-to-gql-fragment';
 import runAction from '../run-action';
 import parseData from '../parse-data';
+import generativeAI from '../generative-ai';
 
 global.generatePDF = generatePDF;
 global.gql = gql;
 global.parseToGqlFragment = parseToGqlFragment;
 global.runAction = runAction;
 global.parseData = parseData;
+global.generativeAI = generativeAI;