pallet_migrations/
mock_helpers.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Test helpers for internal and external usage.
19
20#![allow(missing_docs)]
21
22use alloc::{vec, vec::Vec};
23use codec::{Decode, Encode};
24use frame_support::{
25	migrations::*,
26	weights::{Weight, WeightMeter},
27};
28use sp_core::ConstU32;
29use sp_runtime::BoundedVec;
30
31/// Opaque identifier of a migration.
32pub type MockedIdentifier = BoundedVec<u8, ConstU32<256>>;
33
34/// How a mocked migration should behave.
35#[derive(Debug, Clone, Copy, Encode, Decode)]
36pub enum MockedMigrationKind {
37	/// Succeed after its number of steps elapsed.
38	SucceedAfter,
39	/// Fail after its number of steps elapsed.
40	FailAfter,
41	/// Never terminate.
42	TimeoutAfter,
43	/// Cause an [`SteppedMigrationError::InsufficientWeight`] error after its number of steps
44	/// elapsed.
45	HighWeightAfter(Weight),
46	/// PreUpgrade should fail.
47	#[cfg(feature = "try-runtime")]
48	PreUpgradeFail,
49	/// PostUpgrade should fail.
50	#[cfg(feature = "try-runtime")]
51	PostUpgradeFail,
52}
53use MockedMigrationKind::*; // C style
54
55/// Creates a migration identifier with a specific `kind` and `steps`.
56pub fn mocked_id(kind: MockedMigrationKind, steps: u32) -> MockedIdentifier {
57	(b"MockedMigration", kind, steps).encode().try_into().unwrap()
58}
59
60frame_support::parameter_types! {
61	/// The configs for the migrations to run.
62	storage MIGRATIONS: Vec<(MockedMigrationKind, u32)> = vec![];
63}
64
65/// Allows to set the migrations to run at runtime instead of compile-time.
66///
67/// It achieves this by using the storage to store the migrations to run.
68pub struct MockedMigrations;
69impl SteppedMigrations for MockedMigrations {
70	fn len() -> u32 {
71		MIGRATIONS::get().len() as u32
72	}
73
74	fn nth_id(n: u32) -> Option<Vec<u8>> {
75		let k = MIGRATIONS::get().get(n as usize).copied();
76		k.map(|(kind, steps)| mocked_id(kind, steps).into_inner())
77	}
78
79	fn nth_step(
80		n: u32,
81		cursor: Option<Vec<u8>>,
82		_meter: &mut WeightMeter,
83	) -> Option<Result<Option<Vec<u8>>, SteppedMigrationError>> {
84		let (kind, steps) = MIGRATIONS::get()[n as usize];
85
86		let mut count: u32 =
87			cursor.as_ref().and_then(|c| Decode::decode(&mut &c[..]).ok()).unwrap_or(0);
88		log::debug!("MockedMigration: Step {count} vs max {steps}");
89		if count != steps || matches!(kind, TimeoutAfter) {
90			count += 1;
91			return Some(Ok(Some(count.encode())));
92		}
93
94		Some(match kind {
95			SucceedAfter => {
96				log::debug!("MockedMigration: Succeeded after {count} steps");
97				Ok(None)
98			},
99			HighWeightAfter(required) => {
100				log::debug!("MockedMigration: Not enough weight after {count} steps");
101				Err(SteppedMigrationError::InsufficientWeight { required })
102			},
103			FailAfter => {
104				log::debug!("MockedMigration: Failed after {count} steps");
105				Err(SteppedMigrationError::Failed)
106			},
107			TimeoutAfter => unreachable!(),
108			#[cfg(feature = "try-runtime")]
109			PreUpgradeFail | PostUpgradeFail => Ok(None),
110		})
111	}
112
113	fn nth_transactional_step(
114		n: u32,
115		cursor: Option<Vec<u8>>,
116		meter: &mut WeightMeter,
117	) -> Option<Result<Option<Vec<u8>>, SteppedMigrationError>> {
118		// This is a hack but should be fine. We don't need it in testing.
119		Self::nth_step(n, cursor, meter)
120	}
121
122	fn nth_max_steps(n: u32) -> Option<Option<u32>> {
123		MIGRATIONS::get().get(n as usize).map(|(_, s)| Some(*s))
124	}
125
126	fn nth_migrating_prefixes(n: u32) -> Option<Option<Vec<Vec<u8>>>> {
127		MIGRATIONS::get()
128			.get(n as usize)
129			.map(|(k, s)| Some(vec![mocked_id(*k, *s).into_inner()]))
130	}
131
132	#[cfg(feature = "try-runtime")]
133	fn nth_pre_upgrade(n: u32) -> Option<Result<Vec<u8>, sp_runtime::TryRuntimeError>> {
134		let (kind, _) = MIGRATIONS::get()[n as usize];
135
136		if let PreUpgradeFail = kind {
137			return Some(Err("Some pre-upgrade error".into()));
138		}
139
140		Some(Ok(vec![]))
141	}
142
143	#[cfg(feature = "try-runtime")]
144	fn nth_post_upgrade(
145		n: u32,
146		_state: Vec<u8>,
147	) -> Option<Result<(), sp_runtime::TryRuntimeError>> {
148		let (kind, _) = MIGRATIONS::get()[n as usize];
149
150		if let PostUpgradeFail = kind {
151			return Some(Err("Some post-upgrade error".into()));
152		}
153
154		Some(Ok(()))
155	}
156
157	fn cursor_max_encoded_len() -> usize {
158		65_536
159	}
160
161	fn identifier_max_encoded_len() -> usize {
162		256
163	}
164}
165
166impl MockedMigrations {
167	/// Set the migrations to run.
168	pub fn set(migrations: Vec<(MockedMigrationKind, u32)>) {
169		MIGRATIONS::set(&migrations);
170	}
171}
172
173impl crate::MockedMigrations for MockedMigrations {
174	fn set_fail_after(steps: u32) {
175		MIGRATIONS::set(&vec![(FailAfter, steps)]);
176	}
177
178	fn set_success_after(steps: u32) {
179		MIGRATIONS::set(&vec![(SucceedAfter, steps)]);
180	}
181}
pallet_migrations/mock_helpers.rs

pallet_migrations/
mock_helpers.rs
