Test: Added PostUpgradeArb

.github/workflows/test.yml

@@ -149,6 +149,7 @@ jobs:
                     "property_tests/tests/nat64",
                     "property_tests/tests/null",
                     "property_tests/tests/opt",
+                    "property_tests/tests/post_upgrade_method",
                     "property_tests/tests/principal",
                     "property_tests/tests/query_methods",
                     "property_tests/tests/record",

property_tests/arbitraries/canister_arb.ts

@@ -1,22 +1,36 @@
 import fc from 'fast-check';
 import { QueryMethod } from './canister_methods/query_method_arb';
 import { Test } from '../../test';
-import { UpdateMethod } from './canister_methods/update_method_arb';
-import { InitMethod } from './canister_methods/init_method_arb';
 import { CorrespondingJSType } from './candid/corresponding_js_type';
+import { InitMethod } from './canister_methods/init_method_arb';
+import { PostUpgradeMethod } from './canister_methods/post_upgrade_arb';
+import { UpdateMethod } from './canister_methods/update_method_arb';
 
 export type Canister = {
-    initArgs: string[] | undefined;
+    deployArgs: string[] | undefined;
     sourceCode: string;
     tests: Test[][];
 };
 
+export type CanisterMethod<
+    ParamAgentArgumentValue extends CorrespondingJSType,
+    ParamAgentResponseValue
+> =
+    | QueryMethod
+    | UpdateMethod
+    | InitMethod<ParamAgentArgumentValue, ParamAgentResponseValue>
+    | PostUpgradeMethod<ParamAgentArgumentValue, ParamAgentResponseValue>;
+
 export type CanisterConfig<
     ParamAgentArgumentValue extends CorrespondingJSType = undefined,
     ParamAgentResponseValue = undefined
 > = {
     globalDeclarations?: string[];
     initMethod?: InitMethod<ParamAgentArgumentValue, ParamAgentResponseValue>;
+    postUpgradeMethod?: PostUpgradeMethod<
+        ParamAgentArgumentValue,
+        ParamAgentResponseValue
+    >;
     queryMethods?: QueryMethod[];
     updateMethods?: UpdateMethod[];
 };
@@ -31,17 +45,17 @@ export function CanisterArb<
     >
 ) {
     return configArb.map((config): Canister => {
-        const canisterMethods: (
-            | QueryMethod
-            | UpdateMethod
-            | InitMethod<ParamAgentArgumentValue, ParamAgentResponseValue>
-        )[] = [
+        const canisterMethods: CanisterMethod<
+            ParamAgentArgumentValue,
+            ParamAgentResponseValue
+        >[] = [
             ...(config.initMethod ? [config.initMethod] : []),
+            ...(config.postUpgradeMethod ? [config.postUpgradeMethod] : []),
             ...(config.queryMethods ?? []),
             ...(config.updateMethods ?? [])
         ];
 
-        const initArgs = config.initMethod?.params.map(
+        const deployArgs = config.initMethod?.params.map(
             ({
                 el: {
                     src: { candidTypeAnnotation },
@@ -87,7 +101,7 @@ export function CanisterArb<
         );
 
         return {
-            initArgs,
+            deployArgs,
             sourceCode,
             tests
         };
@@ -99,14 +113,10 @@ function generateSourceCode(
     canisterMethods: (UpdateMethod | QueryMethod)[]
 ) {
     const imports = [
-        ...new Set(
-            canisterMethods.reduce(
-                (acc, method) => {
-                    return [...acc, ...method.imports];
-                },
-                ['Canister', 'query', 'update']
-            )
-        )
+        ...new Set([
+            'Canister',
+            ...canisterMethods.flatMap((method) => [...method.imports])
+        ])
     ]
         .sort((a, b) => a.localeCompare(b, 'en', { sensitivity: 'base' }))
         .join();
@@ -116,8 +126,7 @@ function generateSourceCode(
     );
 
     const declarations = [
-        ...globalDeclarations,
-        ...declarationsFromCanisterMethods
+        ...new Set([...globalDeclarations, ...declarationsFromCanisterMethods])
     ].join('\n');
 
     const sourceCodes = canisterMethods.map((method) => method.sourceCode);

property_tests/arbitraries/canister_methods/post_upgrade_arb.ts

@@ -0,0 +1,141 @@
+import fc from 'fast-check';
+
+import { CandidValueAndMeta } from '../candid/candid_value_and_meta_arb';
+import { CorrespondingJSType } from '../candid/corresponding_js_type';
+import { UniqueIdentifierArb } from '../unique_identifier_arb';
+import { Named } from '../..';
+import {
+    BodyGenerator,
+    TestsGenerator,
+    CallbackLocation,
+    isDefined,
+    generateCallback
+} from '.';
+import { Test } from '../../../test';
+import { VoidArb } from '../candid/primitive/void';
+
+export type PostUpgradeMethod<
+    ParamAgentArgumentValue extends CorrespondingJSType,
+    ParamAgentResponseValue
+> = {
+    imports: Set<string>;
+    globalDeclarations: string[];
+    params: Named<
+        CandidValueAndMeta<ParamAgentArgumentValue, ParamAgentResponseValue>
+    >[];
+    sourceCode: string;
+    tests: Test[][];
+};
+
+export function PostUpgradeMethodArb<
+    ParamAgentArgumentValue extends CorrespondingJSType,
+    ParamAgentResponseValue
+>(
+    paramTypeArrayArb: fc.Arbitrary<
+        CandidValueAndMeta<ParamAgentArgumentValue, ParamAgentResponseValue>[]
+    >,
+    constraints: {
+        generateBody: BodyGenerator<
+            ParamAgentArgumentValue,
+            ParamAgentResponseValue
+        >;
+        generateTests: TestsGenerator<
+            ParamAgentArgumentValue,
+            ParamAgentResponseValue
+        >;
+        callbackLocation?: CallbackLocation;
+    }
+) {
+    return fc
+        .tuple(
+            UniqueIdentifierArb('canisterMethod'),
+            paramTypeArrayArb,
+            VoidArb(),
+            fc.constantFrom(
+                'INLINE',
+                'STANDALONE'
+            ) as fc.Arbitrary<CallbackLocation>,
+            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,
+                defaultCallbackLocation,
+                callbackName
+            ]) => {
+                // TODO: Add a return type to this map callback of type PostUpgradeMethod<Something, Something>
+
+                const callbackLocation =
+                    constraints.callbackLocation ?? defaultCallbackLocation;
+
+                const imports = new Set([
+                    'postUpgrade',
+                    ...paramTypes.flatMap((param) => [...param.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 variableAliasDeclarations = paramTypes
+                    .flatMap((param) => param.src.variableAliasDeclarations)
+                    .filter(isDefined);
+
+                const globalDeclarations =
+                    callbackLocation === 'STANDALONE'
+                        ? [...variableAliasDeclarations, callback]
+                        : variableAliasDeclarations;
+
+                const sourceCode = generateSourceCode(
+                    functionName,
+                    paramTypes,
+                    callbackLocation === 'STANDALONE' ? callbackName : callback
+                );
+
+                const tests = constraints.generateTests(
+                    functionName,
+                    namedParams,
+                    returnType
+                );
+
+                return {
+                    imports,
+                    globalDeclarations,
+                    params: namedParams,
+                    sourceCode,
+                    tests
+                };
+            }
+        );
+}
+
+function generateSourceCode<
+    ParamType extends CorrespondingJSType,
+    ParamAgentType
+>(
+    functionName: string,
+    paramTypes: CandidValueAndMeta<ParamType, ParamAgentType>[],
+    callback: string
+): string {
+    const paramCandidTypeObjects = paramTypes
+        .map((param) => param.src.candidTypeObject)
+        .join(', ');
+
+    return `${functionName}: postUpgrade([${paramCandidTypeObjects}], ${callback})`;
+}

property_tests/arbitraries/canister_methods/query_method_arb.ts

@@ -78,6 +78,7 @@ export function QueryMethodArb<
                     constraints.callbackLocation ?? defaultCallbackLocation;
 
                 const imports = new Set([
+                    'query',
                     ...paramTypes.flatMap((param) => [...param.src.imports]),
                     ...returnType.src.imports
                 ]);

property_tests/arbitraries/canister_methods/update_method_arb.ts

@@ -79,6 +79,7 @@ export function UpdateMethodArb<
                     constraints.callbackLocation ?? defaultCallbackLocation;
 
                 const imports = new Set([
+                    'update',
                     ...paramTypes.flatMap((param) => [...param.src.imports]),
                     ...returnType.src.imports
                 ]);

property_tests/index.ts

@@ -45,8 +45,8 @@ export function runPropTests(canisterArb: fc.Arbitrary<Canister>) {
 
             for (let i = 0; i < canister.tests.length; i++) {
                 const argumentsString =
-                    canister.initArgs && canister.initArgs.length > 0
-                        ? `--argument '(${canister.initArgs.join(', ')})'`
+                    canister.deployArgs && canister.deployArgs.length > 0
+                        ? `--argument '(${canister.deployArgs.join(', ')})'`
                         : '';
 
                 execSync(`dfx deploy canister ${argumentsString}`, {

property_tests/tests/init_method/test/test.ts

@@ -96,7 +96,7 @@ function generateGetInitValuesCanisterMethod(
     globalInitVariableNames: string[]
 ): QueryMethod {
     return {
-        imports: new Set(['Tuple', 'bool']),
+        imports: new Set(['bool', 'query', 'Tuple']),
         globalDeclarations: [],
         sourceCode: /*TS*/ `getInitValues: query(
             [],

property_tests/tests/post_upgrade_method/dfx.json

@@ -0,0 +1,16 @@
+{
+    "canisters": {
+        "canister": {
+            "type": "custom",
+            "main": "src/index.ts",
+            "candid": "src/index.did",
+            "build": "npx azle canister",
+            "wasm": ".azle/canister/canister.wasm",
+            "gzip": true,
+            "declarations": {
+                "output": "test/dfx_generated/canister",
+                "node_compatibility": true
+            }
+        }
+    }
+}