pezkuwi-sdk/bizinikiwi/pezframe/examples/single-block-migrations/src/migrations/v1.rs

// This file is part of Bizinikiwi.

// Copyright (C) Parity Technologies (UK) Ltd. and Dijital Kurdistan Tech Institute
// SPDX-License-Identifier: MIT-0

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
// of the Software, and to permit persons to whom the Software is furnished to do
// so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

use pezframe_support::{
	storage_alias,
	traits::{Get, UncheckedOnRuntimeUpgrade},
};

#[cfg(feature = "try-runtime")]
use alloc::vec::Vec;

/// Collection of storage item formats from the previous storage version.
///
/// Required so we can read values in the v0 storage format during the migration.
mod v0 {
	use super::*;

	/// V0 type for [`crate::Value`].
	#[storage_alias]
	pub type Value<T: crate::Config> = StorageValue<crate::Pezpallet<T>, u32>;
}

/// Implements [`UncheckedOnRuntimeUpgrade`], migrating the state of this pezpallet from V0 to V1.
///
/// In V0 of the template [`crate::Value`] is just a `u32`. In V1, it has been upgraded to
/// contain the struct [`crate::CurrentAndPreviousValue`].
///
/// In this migration, update the on-chain storage for the pezpallet to reflect the new storage
/// layout.
pub struct InnerMigrateV0ToV1<T: crate::Config>(core::marker::PhantomData<T>);

impl<T: crate::Config> UncheckedOnRuntimeUpgrade for InnerMigrateV0ToV1<T> {
	/// Return the existing [`crate::Value`] so we can check that it was correctly set in
	/// `InnerMigrateV0ToV1::post_upgrade`.
	#[cfg(feature = "try-runtime")]
	fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
		use codec::Encode;

		// Access the old value using the `storage_alias` type
		let old_value = v0::Value::<T>::get();
		// Return it as an encoded `Vec<u8>`
		Ok(old_value.encode())
	}

	/// Migrate the storage from V0 to V1.
	///
	/// - If the value doesn't exist, there is nothing to do.
	/// - If the value exists, it is read and then written back to storage inside a
	///   [`crate::CurrentAndPreviousValue`].
	fn on_runtime_upgrade() -> pezframe_support::weights::Weight {
		// Read the old value from storage
		if let Some(old_value) = v0::Value::<T>::take() {
			// Write the new value to storage
			let new = crate::CurrentAndPreviousValue { current: old_value, previous: None };
			crate::Value::<T>::put(new);
			// One read + write for taking the old value, and one write for setting the new value
			T::DbWeight::get().reads_writes(1, 2)
		} else {
			// No writes since there was no old value, just one read for checking
			T::DbWeight::get().reads(1)
		}
	}

	/// Verifies the storage was migrated correctly.
	///
	/// - If there was no old value, the new value should not be set.
	/// - If there was an old value, the new value should be a [`crate::CurrentAndPreviousValue`].
	#[cfg(feature = "try-runtime")]
	fn post_upgrade(state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
		use codec::Decode;
		use pezframe_support::ensure;

		let maybe_old_value = Option::<u32>::decode(&mut &state[..]).map_err(|_| {
			pezsp_runtime::TryRuntimeError::Other("Failed to decode old value from storage")
		})?;

		match maybe_old_value {
			Some(old_value) => {
				let expected_new_value =
					crate::CurrentAndPreviousValue { current: old_value, previous: None };
				let actual_new_value = crate::Value::<T>::get();

				ensure!(actual_new_value.is_some(), "New value not set");
				ensure!(
					actual_new_value == Some(expected_new_value),
					"New value not set correctly"
				);
			},
			None => {
				ensure!(crate::Value::<T>::get().is_none(), "New value unexpectedly set");
			},
		};
		Ok(())
	}
}

/// [`UncheckedOnRuntimeUpgrade`] implementation [`InnerMigrateV0ToV1`] wrapped in a
/// [`VersionedMigration`](pezframe_support::migrations::VersionedMigration), which ensures that:
/// - The migration only runs once when the on-chain storage version is 0
/// - The on-chain storage version is updated to `1` after the migration executes
/// - Reads/Writes from checking/settings the on-chain storage version are accounted for
pub type MigrateV0ToV1<T> = pezframe_support::migrations::VersionedMigration<
	0, // The migration will only execute when the on-chain storage version is 0
	1, // The on-chain storage version will be set to 1 after the migration is complete
	InnerMigrateV0ToV1<T>,
	crate::pezpallet::Pezpallet<T>,
	<T as pezframe_system::Config>::DbWeight,
>;

/// Tests for our migration.
///
/// When writing migration tests, it is important to check:
/// 1. `on_runtime_upgrade` returns the expected weight
/// 2. `post_upgrade` succeeds when given the bytes returned by `pre_upgrade`
/// 3. The storage is in the expected state after the migration
#[cfg(any(all(feature = "try-runtime", test), doc))]
mod test {
	use self::InnerMigrateV0ToV1;
	use super::*;
	use crate::mock::{new_test_ext, MockRuntime};
	use pezframe_support::assert_ok;

	#[test]
	fn handles_no_existing_value() {
		new_test_ext().execute_with(|| {
			// By default, no value should be set. Verify this assumption.
			assert!(crate::Value::<MockRuntime>::get().is_none());
			assert!(v0::Value::<MockRuntime>::get().is_none());

			// Get the pre_upgrade bytes
			let bytes = match InnerMigrateV0ToV1::<MockRuntime>::pre_upgrade() {
				Ok(bytes) => bytes,
				Err(e) => panic!("pre_upgrade failed: {e:?}"),
			};

			// Execute the migration
			let weight = InnerMigrateV0ToV1::<MockRuntime>::on_runtime_upgrade();

			// Verify post_upgrade succeeds
			assert_ok!(InnerMigrateV0ToV1::<MockRuntime>::post_upgrade(bytes));

			// The weight should be just 1 read for trying to access the old value.
			assert_eq!(weight, <MockRuntime as pezframe_system::Config>::DbWeight::get().reads(1));

			// After the migration, no value should have been set.
			assert!(crate::Value::<MockRuntime>::get().is_none());
		})
	}

	#[test]
	fn handles_existing_value() {
		new_test_ext().execute_with(|| {
			// Set up an initial value
			let initial_value = 42;
			v0::Value::<MockRuntime>::put(initial_value);

			// Get the pre_upgrade bytes
			let bytes = match InnerMigrateV0ToV1::<MockRuntime>::pre_upgrade() {
				Ok(bytes) => bytes,
				Err(e) => panic!("pre_upgrade failed: {e:?}"),
			};

			// Execute the migration
			let weight = InnerMigrateV0ToV1::<MockRuntime>::on_runtime_upgrade();

			// Verify post_upgrade succeeds
			assert_ok!(InnerMigrateV0ToV1::<MockRuntime>::post_upgrade(bytes));

			// The weight used should be 1 read for the old value, and 1 write for the new
			// value.
			assert_eq!(
				weight,
				<MockRuntime as pezframe_system::Config>::DbWeight::get().reads_writes(1, 2)
			);

			// After the migration, the new value should be set as the `current` value.
			let expected_new_value =
				crate::CurrentAndPreviousValue { current: initial_value, previous: None };
			assert_eq!(crate::Value::<MockRuntime>::get(), Some(expected_new_value));
		})
	}
}