Add CallbackLocation

demergent-labs · Dec 1, 2023 · fb6a5f9 · fb6a5f9
1 parent 98b54e7
commit fb6a5f9
Show file tree

Hide file tree

Showing 4 changed files with 159 additions and 65 deletions.
diff --git a/property_tests/arbitraries/canister_arb.ts b/property_tests/arbitraries/canister_arb.ts
@@ -39,11 +39,8 @@ function generateSourceCode(queryMethods: QueryMethod[]) {
         )
     ].join();
 
-    const candidTypeDeclarations = queryMethods
-        .map(
-            (queryMethod) =>
-                queryMethod.candidTypeDeclarations?.join('\n') ?? ''
-        )
+    const declarations = queryMethods
+        .flatMap((queryMethod) => queryMethod.globalDeclarations)
         .join('\n');
 
     const sourceCodes = queryMethods.map(
@@ -56,7 +53,7 @@ function generateSourceCode(queryMethods: QueryMethod[]) {
         // TODO solve the underlying principal problem https://github.com/demergent-labs/azle/issues/1443
         import { Principal as DfinityPrincipal } from '@dfinity/principal';
 
-        ${candidTypeDeclarations}
+        ${declarations}
 
         export default Canister({
             ${sourceCodes.join(',\n    ')}

diff --git a/property_tests/arbitraries/query_method_arb.ts b/property_tests/arbitraries/query_method_arb.ts
@@ -1,13 +1,15 @@
 import fc from 'fast-check';
-import { Test } from '../../test';
-import { UniqueIdentifierArb } from './unique_identifier_arb';
-import { CandidType } from './candid/candid_type_arb';
+
 import { CandidMeta } from './candid/candid_arb';
+import { CandidReturnType } from './candid/candid_return_type_arb';
+import { CandidType } from './candid/candid_type_arb';
+import { UniqueIdentifierArb } from './unique_identifier_arb';
+import { Test } from '../../test';
 import { Named } from '../';
 
 export type QueryMethod = {
     imports: Set<string>;
-    candidTypeDeclarations: string[] | undefined;
+    globalDeclarations: string[];
     sourceCode: string;
     tests: Test[];
 };
@@ -43,6 +45,11 @@ export type TestsGenerator<
     >
 ) => Test[];
 
+export enum CallbackLocation {
+    Inline = 'INLINE',
+    Standalone = 'STANDALONE'
+}
+
 export function QueryMethodArb<
     ParamAgentArgumentValue extends CandidType,
     ParamAgentResponseValue,
@@ -68,58 +75,143 @@ export function QueryMethodArb<
             ReturnTypeAgentArgumentValue,
             ReturnTypeAgentResponseValue
         >;
-        // TODO: Consider adding a callback to determine the returnType
-        // i.e. instead of using the first one if the params array isn't empty.
+        callbackLocation?: CallbackLocation;
     }
 ) {
     return fc
         .tuple(
             UniqueIdentifierArb('canisterMethod'),
             paramTypeArrayArb,
-            returnTypeArb
+            returnTypeArb,
+            fc.constantFrom(
+                CallbackLocation.Inline,
+                CallbackLocation.Standalone
+            ),
+            UniqueIdentifierArb('typeDeclaration')
+            // TODO: This unique id would be better named globalScope or something
+            // But needs to match the same scope as typeDeclarations so I'm using
+            // that for now.
         )
-        .map(([functionName, paramTypes, returnType]): QueryMethod => {
-            const imports = new Set([
-                ...paramTypes.flatMap((param) => [...param.src.imports]),
-                ...returnType.src.imports
-            ]);
-
-            const candidTypeDeclarations = [
-                ...paramTypes.map((param) => param.src.typeDeclaration ?? ''),
-                returnType.src.typeDeclaration ?? ''
-            ];
-
-            const namedParams = paramTypes.map(
-                <T>(param: T, index: number): Named<T> => ({
-                    name: `param${index}`,
-                    el: param
-                })
-            );
-
-            const paramCandidTypes = paramTypes
-                .map((param) => param.src.candidType)
-                .join(', ');
-
-            const returnCandidType = returnType.src.candidType;
-
-            const paramNames = namedParams
-                .map((namedParam) => namedParam.name)
-                .join(', ');
-
-            const body = constraints.generateBody(namedParams, returnType);
-            const tests = constraints.generateTests(
+        .map(
+            ([
                 functionName,
-                namedParams,
-                returnType
-            );
-
-            return {
-                imports,
-                candidTypeDeclarations,
-                sourceCode: `${functionName}: query([${paramCandidTypes}], ${returnCandidType}, (${paramNames}) => {
-                    ${body}
-                })`,
-                tests
-            };
-        });
+                paramTypes,
+                returnType,
+                defaultCallbackLocation,
+                callbackName
+            ]): QueryMethod => {
+                const callbackLocation =
+                    constraints.callbackLocation ?? defaultCallbackLocation;
+
+                const imports = new Set([
+                    ...paramTypes.flatMap((param) => [...param.src.imports]),
+                    ...returnType.src.imports
+                ]);
+
+                const namedParams = paramTypes.map(
+                    <T>(param: T, index: number): Named<T> => ({
+                        name: `param${index}`,
+                        el: param
+                    })
+                );
+
+                const callback = generateCallback(
+                    namedParams,
+                    returnType,
+                    constraints.generateBody,
+                    callbackLocation,
+                    callbackName
+                );
+
+                const candidTypeDeclarations = [
+                    ...paramTypes.map((param) => param.src.typeDeclaration),
+                    returnType.src.typeDeclaration
+                ].filter(isDefined);
+
+                const globalDeclarations =
+                    callbackLocation === CallbackLocation.Standalone
+                        ? [...candidTypeDeclarations, callback]
+                        : candidTypeDeclarations;
+
+                const sourceCode = generateSourceCode(
+                    functionName,
+                    paramTypes,
+                    returnType,
+                    callbackLocation === CallbackLocation.Standalone
+                        ? callbackName
+                        : callback
+                );
+
+                const tests = constraints.generateTests(
+                    functionName,
+                    namedParams,
+                    returnType
+                );
+
+                return {
+                    imports,
+                    globalDeclarations,
+                    sourceCode,
+                    tests
+                };
+            }
+        );
+}
+
+function isDefined<T>(value: T | undefined): value is T {
+    return value !== undefined;
+}
+
+function generateCallback<
+    ParamType extends CandidType,
+    ParamAgentType,
+    ReturnType extends CandidReturnType,
+    ReturnAgentType
+>(
+    namedParams: Named<CandidMeta<ParamType, ParamAgentType>>[],
+    returnType: CandidMeta<ReturnType, ReturnAgentType>,
+    generateBody: BodyGenerator<
+        ParamType,
+        ParamAgentType,
+        ReturnType,
+        ReturnAgentType
+    >,
+    callbackLocation: CallbackLocation,
+    callbackName: string
+): string {
+    const paramNames = namedParams
+        .map((namedParam) => namedParam.name)
+        .join(', ');
+
+    const body = generateBody(namedParams, returnType);
+
+    if (callbackLocation === CallbackLocation.Inline) {
+        return `(${paramNames}) => {${body}}`;
+    }
+
+    const paramNamesAndTypes = namedParams
+        .map((namedParam) => `${namedParam.name}: any`) // TODO: Use actual candid type, not any
+        .join(', ');
+
+    return `function ${callbackName}(${paramNamesAndTypes}) {${body}}`;
+}
+
+function generateSourceCode<
+    ParamType extends CandidType,
+    ParamAgentType,
+    ReturnType extends CandidReturnType,
+    ReturnAgentType
+>(
+    functionName: string,
+    paramTypes: CandidMeta<ParamType, ParamAgentType>[],
+    returnType: CandidMeta<ReturnType, ReturnAgentType>,
+    callback: string
+): string {
+    const paramCandidTypes = paramTypes
+        .map((param) => param.src.candidType)
+        .join(', ');
+
+    const returnCandidType = returnType.src.candidType;
+
+    return `${functionName}: query([${paramCandidTypes}], ${returnCandidType}, ${callback})`;
 }
diff --git a/property_tests/tests/query_methods/test/generate_body.ts b/property_tests/tests/query_methods/test/generate_body.ts
@@ -1,11 +1,14 @@
 import { areParamsCorrectlyOrdered } from 'azle/property_tests/are_params_correctly_ordered';
 import { CandidType } from 'azle/property_tests/arbitraries/candid/candid_type_arb';
 import { BodyGenerator } from 'azle/property_tests/arbitraries/query_method_arb';
+import { CandidReturnType } from 'azle/property_tests/arbitraries/candid/candid_return_type_arb';
 
-export const generateBody: BodyGenerator<CandidType, CandidType> = (
-    namedParams,
-    returnType
-): string => {
+export const generateBody: BodyGenerator<
+    CandidType,
+    CandidReturnType,
+    CandidType,
+    CandidReturnType
+> = (namedParams, returnType): string => {
     const paramsAreCorrectlyOrdered = areParamsCorrectlyOrdered(namedParams);
 
     return `

diff --git a/property_tests/tests/query_methods/test/generate_tests.ts b/property_tests/tests/query_methods/test/generate_tests.ts
@@ -4,12 +4,14 @@ import { getActor } from 'azle/property_tests';
 import { CandidType } from 'azle/property_tests/arbitraries/candid/candid_type_arb';
 import { TestsGenerator } from 'azle/property_tests/arbitraries/query_method_arb';
 import { Test } from 'azle/test';
+import { CandidReturnType } from '../../../arbitraries/candid/candid_return_type_arb';
 
-export const generateTests: TestsGenerator<CandidType, CandidType> = (
-    functionName,
-    params,
-    returnType
-): Test[] => {
+export const generateTests: TestsGenerator<
+    CandidType,
+    CandidReturnType,
+    CandidType,
+    CandidReturnType
+> = (functionName, params, returnType): Test[] => {
     const paramValues = params.map((param) => param.el.agentArgumentValue);
     const expectedResult = returnType.agentResponseValue;