From cd2a1cb50862601d5a3d02eb82afe996770a5f7f Mon Sep 17 00:00:00 2001
From: Max Andreev <maksandre@yandex.ru>
Date: Fri, 9 Jun 2023 14:04:00 +0000
Subject: [PATCH] add before and after hooks to frankenstein

---
 .env                                  | 1 +
 tests/src/util/frankenstein.ts        | 5 +++++
 tests/src/util/frankensteinMigrate.ts | 9 +++++++++
 3 files changed, 15 insertions(+)
 create mode 100644 tests/src/util/frankensteinMigrate.ts

diff --git a/.env b/.env
index c3a9b1c725..10976c2c42 100644
--- a/.env
+++ b/.env
@@ -2,6 +2,7 @@ RUST_TOOLCHAIN=nightly-2022-11-15
 POLKADOT_LAUNCH_BRANCH=unique-network
 RELAY_CHAIN_TYPE=westend
 CHAINQL=v0.4.1
+DESTINATION_SPEC_VERSION=v942057
 
 POLKADOT_MAINNET_BRANCH=release-v0.9.37
 STATEMINT_BUILD_BRANCH=release-parachains-v9370
diff --git a/tests/src/util/frankenstein.ts b/tests/src/util/frankenstein.ts
index 0babf59cb1..8d6da267f4 100644
--- a/tests/src/util/frankenstein.ts
+++ b/tests/src/util/frankenstein.ts
@@ -20,6 +20,7 @@ import zombie from '@zombienet/orchestrator/dist';
 import {readNetworkConfig} from '@zombienet/utils/dist';
 import {resolve} from 'path';
 import {usingPlaygrounds} from '.';
+import {migrations} from './frankensteinMigrate';
 import fs from 'fs';
 
 const ZOMBIENET_CREDENTIALS = process.env.ZOMBIENET_CREDENTIALS || '../.env';
@@ -245,12 +246,14 @@ const raiseZombienet = async (): Promise<void> => {
       await waitWithTimer(relayInfo.epochTime);
     }
 
+    const migration = migrations[process.env.DESTINATION_SPEC_VERSION!];
     for(const paraId in network.paras) {
       console.log(`\n--- Upgrading the runtime of parachain ${paraId} \t---`);
       const para = network.paras[paraId];
 
       // Enable maintenance mode if present
       await toggleMaintenanceMode(true, para.nodes[0].wsUri);
+      if(migration) await migration.before();
 
       // Read the WASM code and authorize the upgrade with its hash and set it as the new runtime
       const code = fs.readFileSync(NEW_PARA_WASM);
@@ -324,6 +327,8 @@ const raiseZombienet = async (): Promise<void> => {
 
     // Disable maintenance mode if present
     for(const paraId in network.paras) {
+      // TODO only if our parachain
+      if(migration) await migration.after();
       await toggleMaintenanceMode(false, network.paras[paraId].nodes[0].wsUri);
     }
   } else {
diff --git a/tests/src/util/frankensteinMigrate.ts b/tests/src/util/frankensteinMigrate.ts
new file mode 100644
index 0000000000..bdb61ffb83
--- /dev/null
+++ b/tests/src/util/frankensteinMigrate.ts
@@ -0,0 +1,9 @@
+import {migration as locksToFreezesMigration} from '../migrations/942057-appPromotion';
+export interface Migration {
+  before: () => Promise<void>,
+  after: () => Promise<void>,
+}
+
+export const migrations: {[key: string]: Migration} = {
+  'v942057': locksToFreezesMigration,
+};