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
    crate::{
19
        assignment::{Assignment, AssignmentError},
20
        tests::Test,
21
    },
22
    rand::{seq::SliceRandom, SeedableRng},
23
    rand_chacha::ChaCha20Rng,
24
    sp_std::collections::btree_map::BTreeMap,
25
};
26

            
27
8
fn no_shuffle() -> Option<fn(&mut Vec<u64>)> {
28
8
    None
29
8
}
30

            
31
#[test]
32
1
fn assign_full_old_assigned_priority() {
33
1
    // Collators in old_assigned will be selected before other collators
34
1
    let collators = vec![1, 2, 3, 4, 5];
35
1
    let container_chains = vec![(1000.into(), 5)];
36
1
    let old_assigned = BTreeMap::from_iter(vec![(1000.into(), vec![3, 4])]);
37
1

            
38
1
    let new_assigned =
39
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
40
1
            .unwrap();
41
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![3, 4, 1, 2, 5])]);
42
1
    assert_eq!(new_assigned, expected);
43
1
}
44

            
45
#[test]
46
1
fn assign_full_invalid_old_assigned_collators_removed() {
47
1
    // If the collators in old_assigned are no longer collators, they are not assigned
48
1
    let collators = vec![1, 2, 3, 4, 5];
49
1
    let container_chains = vec![(1000.into(), 5)];
50
1
    let old_assigned = BTreeMap::from_iter(vec![(1000.into(), vec![20, 21])]);
51
1

            
52
1
    let new_assigned =
53
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
54
1
            .unwrap();
55
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![1, 2, 3, 4, 5])]);
56
1
    assert_eq!(new_assigned, expected);
57
1
}
58

            
59
#[test]
60
1
fn assign_full_invalid_chains_removed() {
61
1
    // Mark all collators as already assigned to a chain that does not exist. Should treat them as not assigned.
62
1
    let collators = vec![1, 2, 3, 4, 5];
63
1
    let container_chains = vec![(1000.into(), 5)];
64
1
    let old_assigned = BTreeMap::from_iter(vec![(1001.into(), vec![1, 2, 3, 4, 5])]);
65
1

            
66
1
    let new_assigned =
67
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
68
1
            .unwrap();
69
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![1, 2, 3, 4, 5])]);
70
1
    assert_eq!(new_assigned, expected);
71
1
}
72

            
73
#[test]
74
1
fn assign_full_truncates_collators() {
75
1
    // Need 2 collators for each chain, when old_assigned has more than 2. Should truncate old_assigned to 2.
76
1
    let collators = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
77
1
    let container_chains = vec![(1000.into(), 2), (2000.into(), 2)];
78
1
    let old_assigned = BTreeMap::from_iter(vec![
79
1
        (1000.into(), vec![1, 2, 3, 4, 5]),
80
1
        (2000.into(), vec![6, 7, 8, 9, 10]),
81
1
    ]);
82
1

            
83
1
    let new_assigned =
84
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
85
1
            .unwrap();
86
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![1, 2]), (2000.into(), vec![6, 7])]);
87
1
    assert_eq!(new_assigned, expected);
88
1
}
89

            
90
#[test]
91
1
fn assign_full_old_assigned_error_if_not_enough_collators() {
92
1
    // Need 4 collators, only have 2, and all 2 were assigned to the second chain. If the function did not panic, we
93
1
    // would have 0 collators assigned to the first chain, which is supposed to have priority.
94
1
    let collators = vec![1, 2];
95
1
    let container_chains = vec![(1000.into(), 2), (2000.into(), 2)];
96
1
    let old_assigned = BTreeMap::from_iter(vec![(2000.into(), vec![1, 2])]);
97
1
    let new_assigned =
98
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle());
99
1
    assert_eq!(
100
1
        new_assigned.unwrap_err(),
101
1
        AssignmentError::NotEnoughCollators
102
1
    );
103
1
}
104

            
105
#[test]
106
1
fn assign_full_list_priority() {
107
1
    // The order in the collators list is priority
108
1
    let collators = vec![
109
1
        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
110
1
    ];
111
1
    let container_chains = vec![(1000.into(), 2), (2000.into(), 2)];
112
1
    let old_assigned = BTreeMap::from_iter(vec![]);
113
1

            
114
1
    let new_assigned =
115
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
116
1
            .unwrap();
117
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![1, 2]), (2000.into(), vec![3, 4])]);
118
1
    assert_eq!(new_assigned, expected);
119
1
}
120

            
121
#[test]
122
1
fn assign_full_list_priority_shuffle() {
123
1
    // The order in the collators list is priority
124
1
    let collators = vec![
125
1
        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
126
1
    ];
127
1
    let container_chains = vec![(1000.into(), 2), (2000.into(), 2)];
128
1
    let old_assigned = BTreeMap::from_iter(vec![]);
129
1
    let shuffle = Some(move |collators: &mut Vec<u64>| {
130
1
        // Shuffle with a fixed seed, we do not need randomness in a unit test
131
1
        let seed = [1; 32];
132
1
        let mut rng: ChaCha20Rng = SeedableRng::from_seed(seed);
133
1
        collators.shuffle(&mut rng);
134
1
    });
135
1

            
136
1
    let new_assigned =
137
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, shuffle)
138
1
            .unwrap();
139
1
    // Expect only [1, 2, 3, 4] to be assigned, in random order
140
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![3, 2]), (2000.into(), vec![1, 4])]);
141
1
    assert_eq!(new_assigned, expected);
142
1
}
143

            
144
#[test]
145
1
fn assign_full_solochain() {
146
1
    // In solochain mode, orchestrator chain has 0 collators. This shouldn't cause any panics.
147
1
    let collators = vec![1, 2, 3, 4, 5];
148
1
    let container_chains = vec![(1000.into(), 0), (2000.into(), 3), (2001.into(), 2)];
149
1
    let old_assigned = BTreeMap::from_iter(vec![]);
150
1

            
151
1
    let new_assigned =
152
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
153
1
            .unwrap();
154
1
    let expected = BTreeMap::from_iter(vec![
155
1
        (1000.into(), vec![]),
156
1
        (2000.into(), vec![1, 2, 3]),
157
1
        (2001.into(), vec![4, 5]),
158
1
    ]);
159
1
    assert_eq!(new_assigned, expected);
160
1
}
161

            
162
#[test]
163
1
fn assign_full_solochain_zero_collators() {
164
1
    // In solochain mode, there can be 0 collators. This shouldn't cause any panics.
165
1
    let collators = vec![];
166
1
    let container_chains = vec![(1000.into(), 0)];
167
1
    let old_assigned = BTreeMap::from_iter(vec![]);
168
1

            
169
1
    let new_assigned =
170
1
        Assignment::<Test>::assign_full(collators, container_chains, old_assigned, no_shuffle())
171
1
            .unwrap();
172
1
    let expected = BTreeMap::from_iter(vec![(1000.into(), vec![])]);
173
1
    assert_eq!(new_assigned, expected);
174
1
}