1
// Copyright (C) Moondance Labs Ltd.
2
// This file is part of Tanssi.
3

            
4
// Tanssi is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Tanssi is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
16

            
17
use {
18
    cumulus_primitives_core::{ParaId, PersistedValidationData},
19
    cumulus_primitives_parachain_inherent::ParachainInherentData,
20
    dancebox_runtime::{AuthorInherent, RuntimeOrigin},
21
    dp_consensus::runtime_decl_for_tanssi_authority_assignment_api::TanssiAuthorityAssignmentApi,
22
    frame_support::{
23
        assert_ok,
24
        traits::{OnFinalize, OnInitialize},
25
    },
26
    nimbus_primitives::{NimbusId, NIMBUS_ENGINE_ID},
27
    pallet_registrar_runtime_api::ContainerChainGenesisData,
28
    parity_scale_codec::{Decode, Encode, MaxEncodedLen},
29
    polkadot_parachain_primitives::primitives::HeadData,
30
    sp_consensus_aura::AURA_ENGINE_ID,
31
    sp_core::Get,
32
    sp_runtime::{traits::Dispatchable, Digest, DigestItem},
33
    sp_std::collections::btree_map::BTreeMap,
34
};
35

            
36
pub use dancebox_runtime::{
37
    AccountId, AssetRate, AuthorNoting, AuthorityAssignment, AuthorityMapping, Balance, Balances,
38
    CollatorAssignment, Configuration, DataPreservers, ForeignAssets, ForeignAssetsCreator,
39
    InflationRewards, Initializer, Invulnerables, MinimumSelfDelegation, ParachainInfo,
40
    PooledStaking, Proxy, ProxyType, Registrar, RewardsPortion, Runtime, RuntimeCall,
41
    ServicesPayment, Session, System, TransactionPayment,
42
};
43

            
44
// TODO: This module copy/pasted for now from dancebox/tests/common/mod.rs, should be extracted and re-used in both places
45

            
46
pub const UNIT: Balance = 1_000_000_000_000;
47

            
48
12
pub fn session_to_block(n: u32) -> u32 {
49
12
    let block_number = dancebox_runtime::Period::get() * n;
50
12

            
51
12
    // Add 1 because the block that emits the NewSession event cannot contain any extrinsics,
52
12
    // so this is the first block of the new session that can actually be used
53
12
    block_number + 1
54
12
}
55

            
56
#[derive(Debug, Clone, Eq, PartialEq)]
57
pub struct RunSummary {
58
    pub author_id: AccountId,
59
    pub inflation: Balance,
60
}
61

            
62
12
pub fn run_to_session(n: u32) {
63
12
    run_to_block(session_to_block(n));
64
12
}
65

            
66
/// Utility function that advances the chain to the desired block number.
67
///
68
/// After this function returns, the current block number will be `n`, and the block will be "open",
69
/// meaning that on_initialize has been executed, but on_finalize has not. To execute on_finalize as
70
/// well, for example to test a runtime api, manually call `end_block` after this, run the test, and
71
/// call `start_block` to ensure that this function keeps working as expected.
72
/// Extrinsics should always be executed before on_finalize.
73
12
pub fn run_to_block(n: u32) -> BTreeMap<u32, RunSummary> {
74
12
    let current_block_number = System::block_number();
75
12
    assert!(
76
12
        current_block_number < n,
77
        "run_to_block called with block {} when current block is {}",
78
        n,
79
        current_block_number
80
    );
81
12
    let mut summaries = BTreeMap::new();
82

            
83
202
    while System::block_number() < n {
84
190
        let summary = run_block();
85
190
        let block_number = System::block_number();
86
190
        summaries.insert(block_number, summary);
87
190
    }
88

            
89
12
    summaries
90
12
}
91

            
92
202
pub fn insert_authorities_and_slot_digests(slot: u64) {
93
202
    let authorities =
94
202
        Runtime::para_id_authorities(ParachainInfo::get()).expect("authorities should be set");
95
202

            
96
202
    let authority: NimbusId = authorities[slot as usize % authorities.len()].clone();
97
202

            
98
202
    let pre_digest = Digest {
99
202
        logs: vec![
100
202
            DigestItem::PreRuntime(AURA_ENGINE_ID, slot.encode()),
101
202
            DigestItem::PreRuntime(NIMBUS_ENGINE_ID, authority.encode()),
102
202
        ],
103
202
    };
104
202

            
105
202
    System::reset_events();
106
202
    System::initialize(
107
202
        &(System::block_number() + 1),
108
202
        &System::parent_hash(),
109
202
        &pre_digest,
110
202
    );
111
202
}
112

            
113
// Used to create the next block inherent data
114
#[derive(Clone, Encode, Decode, Default, PartialEq, Debug, scale_info::TypeInfo, MaxEncodedLen)]
115
pub struct MockInherentData {
116
    pub random_seed: Option<[u8; 32]>,
117
}
118

            
119
202
fn take_new_inherent_data() -> Option<MockInherentData> {
120
202
    let data: Option<MockInherentData> =
121
202
        frame_support::storage::unhashed::take(b"__mock_new_inherent_data");
122
202

            
123
202
    data
124
202
}
125

            
126
#[derive(Clone, Encode, Decode, PartialEq, Debug, scale_info::TypeInfo, MaxEncodedLen)]
127
enum RunBlockState {
128
    Start(u32),
129
    End(u32),
130
}
131

            
132
impl RunBlockState {
133
    fn assert_can_advance(&self, new_state: &RunBlockState) {
134
        match self {
135
            RunBlockState::Start(n) => {
136
                assert_eq!(
137
                    new_state,
138
                    &RunBlockState::End(*n),
139
                    "expected a call to end_block({}), but user called {:?}",
140
                    *n,
141
                    new_state
142
                );
143
            }
144
            RunBlockState::End(n) => {
145
                assert_eq!(
146
                    new_state,
147
                    &RunBlockState::Start(*n + 1),
148
                    "expected a call to start_block({}), but user called {:?}",
149
                    *n + 1,
150
                    new_state
151
                )
152
            }
153
        }
154
    }
155
}
156

            
157
392
fn advance_block_state_machine(new_state: RunBlockState) {
158
392
    if frame_support::storage::unhashed::exists(b"__mock_is_xcm_test") {
159
        // Disable this check in XCM tests, because the XCM emulator runs on_initialize and
160
        // on_finalize automatically
161
392
        return;
162
    }
163
    let old_state: RunBlockState =
164
        frame_support::storage::unhashed::get(b"__mock_debug_block_state").unwrap_or(
165
            // Initial state is expecting a call to start() with block number 1, so old state should be
166
            // end of block 0
167
            RunBlockState::End(0),
168
        );
169
    old_state.assert_can_advance(&new_state);
170
    frame_support::storage::unhashed::put(b"__mock_debug_block_state", &new_state);
171
392
}
172

            
173
202
pub fn start_block() -> RunSummary {
174
202
    let block_number = System::block_number();
175
202
    advance_block_state_machine(RunBlockState::Start(block_number + 1));
176
202
    let mut slot = current_slot() + 1;
177
202
    if block_number == 0 {
178
        // Hack to avoid breaking all tests. When the current block is 1, the slot number should be
179
        // 1. But all of our tests assume it will be 0. So use slot number = block_number - 1.
180
        slot = 0;
181
202
    }
182

            
183
202
    let maybe_mock_inherent = take_new_inherent_data();
184

            
185
202
    if let Some(mock_inherent_data) = maybe_mock_inherent {
186
        set_parachain_inherent_data(mock_inherent_data);
187
202
    }
188

            
189
202
    insert_authorities_and_slot_digests(slot);
190
202

            
191
202
    // Initialize the new block
192
202
    CollatorAssignment::on_initialize(System::block_number());
193
202
    Session::on_initialize(System::block_number());
194
202
    Initializer::on_initialize(System::block_number());
195
202
    AuthorInherent::on_initialize(System::block_number());
196
202

            
197
202
    // `Initializer::on_finalize` needs to run at least one to have
198
202
    // author mapping setup.
199
202
    let author_id = current_author();
200
202

            
201
202
    let current_issuance = Balances::total_issuance();
202
202
    InflationRewards::on_initialize(System::block_number());
203
202
    let new_issuance = Balances::total_issuance();
204
202

            
205
202
    pallet_author_inherent::Pallet::<Runtime>::kick_off_authorship_validation(None.into())
206
202
        .expect("author inherent to dispatch correctly");
207
202

            
208
202
    RunSummary {
209
202
        author_id,
210
202
        inflation: new_issuance - current_issuance,
211
202
    }
212
202
}
213

            
214
190
pub fn end_block() {
215
190
    let block_number = System::block_number();
216
190
    advance_block_state_machine(RunBlockState::End(block_number));
217
190
    // Finalize the block
218
190
    CollatorAssignment::on_finalize(System::block_number());
219
190
    Session::on_finalize(System::block_number());
220
190
    Initializer::on_finalize(System::block_number());
221
190
    AuthorInherent::on_finalize(System::block_number());
222
190
    TransactionPayment::on_finalize(System::block_number());
223
190
}
224

            
225
190
pub fn run_block() -> RunSummary {
226
190
    end_block();
227
190

            
228
190
    start_block()
229
190
}
230

            
231
/// Mock the inherent that sets validation data in ParachainSystem, which
232
/// contains the `relay_chain_block_number`, which is used in `collator-assignment` as a
233
/// source of randomness.
234
pub fn set_parachain_inherent_data(mock_inherent_data: MockInherentData) {
235
    use {
236
        cumulus_primitives_core::relay_chain::well_known_keys,
237
        cumulus_test_relay_sproof_builder::RelayStateSproofBuilder,
238
    };
239

            
240
    let relay_sproof = RelayStateSproofBuilder {
241
        para_id: 100u32.into(),
242
        included_para_head: Some(HeadData(vec![1, 2, 3])),
243
        current_slot: (current_slot()).into(),
244
        additional_key_values: if mock_inherent_data.random_seed.is_some() {
245
            vec![(
246
                well_known_keys::CURRENT_BLOCK_RANDOMNESS.to_vec(),
247
                Some(mock_inherent_data.random_seed).encode(),
248
            )]
249
        } else {
250
            vec![]
251
        },
252
        ..Default::default()
253
    };
254

            
255
    let (relay_parent_storage_root, relay_chain_state) = relay_sproof.into_state_root_and_proof();
256
    let vfp = PersistedValidationData {
257
        relay_parent_number: 1u32,
258
        relay_parent_storage_root,
259
        ..Default::default()
260
    };
261
    let parachain_inherent_data = ParachainInherentData {
262
        validation_data: vfp,
263
        relay_chain_state,
264
        downward_messages: Default::default(),
265
        horizontal_messages: Default::default(),
266
    };
267

            
268
    // Delete existing flag to avoid error
269
    // 'ValidationData must be updated only once in a block'
270
    // TODO: this is a hack
271
    frame_support::storage::unhashed::kill(&frame_support::storage::storage_prefix(
272
        b"ParachainSystem",
273
        b"ValidationData",
274
    ));
275

            
276
    assert_ok!(RuntimeCall::ParachainSystem(
277
        cumulus_pallet_parachain_system::Call::<Runtime>::set_validation_data {
278
            data: parachain_inherent_data
279
        }
280
    )
281
    .dispatch(inherent_origin()));
282
}
283

            
284
pub fn inherent_origin() -> <Runtime as frame_system::Config>::RuntimeOrigin {
285
    <Runtime as frame_system::Config>::RuntimeOrigin::none()
286
}
287

            
288
12
pub fn empty_genesis_data() -> ContainerChainGenesisData {
289
12
    ContainerChainGenesisData {
290
12
        storage: Default::default(),
291
12
        name: Default::default(),
292
12
        id: Default::default(),
293
12
        fork_id: Default::default(),
294
12
        extensions: Default::default(),
295
12
        properties: Default::default(),
296
12
    }
297
12
}
298

            
299
202
pub fn current_slot() -> u64 {
300
202
    u64::from(
301
202
        pallet_async_backing::SlotInfo::<Runtime>::get()
302
202
            .unwrap_or_default()
303
202
            .0,
304
202
    )
305
202
}
306

            
307
202
pub fn current_author() -> AccountId {
308
202
    let current_session = Session::current_index();
309
202
    let mapping =
310
202
        pallet_authority_mapping::Pallet::<Runtime>::authority_id_mapping(current_session)
311
202
            .expect("there is a mapping for the current session");
312
202

            
313
202
    let author = pallet_author_inherent::Author::<Runtime>::get()
314
202
        .expect("there should be a registered author");
315
202

            
316
202
    mapping
317
202
        .get(&author)
318
202
        .expect("there is a mapping for the current author")
319
202
        .clone()
320
202
}
321

            
322
12
pub fn set_dummy_boot_node(para_manager: RuntimeOrigin, para_id: ParaId) {
323
    use {
324
        pallet_data_preservers::{ParaIdsFilter, Profile, ProfileMode},
325
        tp_data_preservers_common::{AssignerExtra, ProviderRequest},
326
    };
327

            
328
12
    let profile = Profile {
329
12
        url:
330
12
            b"/ip4/127.0.0.1/tcp/33049/ws/p2p/12D3KooWHVMhQDHBpj9vQmssgyfspYecgV6e3hH1dQVDUkUbCYC9"
331
12
                .to_vec()
332
12
                .try_into()
333
12
                .expect("to fit in BoundedVec"),
334
12
        para_ids: ParaIdsFilter::AnyParaId,
335
12
        mode: ProfileMode::Bootnode,
336
12
        assignment_request: ProviderRequest::Free,
337
12
    };
338
12

            
339
12
    let profile_id = pallet_data_preservers::NextProfileId::<Runtime>::get();
340
12
    let profile_owner = AccountId::new([1u8; 32]);
341
12
    DataPreservers::force_create_profile(RuntimeOrigin::root(), profile, profile_owner)
342
12
        .expect("profile create to succeed");
343
12

            
344
12
    DataPreservers::start_assignment(para_manager, profile_id, para_id, AssignerExtra::Free)
345
12
        .expect("assignement to work");
346
12

            
347
12
    assert!(
348
12
        pallet_data_preservers::Assignments::<Runtime>::get(para_id).contains(&profile_id),
349
        "profile should be correctly assigned"
350
    );
351
12
}