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, ChainNumCollators},
20
        tests::Test,
21
    },
22
    sp_std::collections::{btree_map::BTreeMap, btree_set::BTreeSet},
23
};
24

            
25
#[test]
26
1
fn invulnerable_priority_0_collators() {
27
1
    let collators = vec![];
28
1
    let orchestrator_chain = ChainNumCollators {
29
1
        para_id: 1000.into(),
30
1
        min_collators: 2,
31
1
        max_collators: 5,
32
1
        parathread: false,
33
1
    };
34
1
    let mut old_assigned = BTreeMap::new();
35
1

            
36
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
37
1
        &collators,
38
1
        orchestrator_chain,
39
1
        &mut old_assigned,
40
1
    );
41
1

            
42
1
    assert_eq!(num_invulnerables, 0);
43
1
}
44

            
45
#[test]
46
1
fn invulnerable_priority_0_invulnerables() {
47
1
    let collators = vec![1, 2, 3, 4, 5];
48
1
    let orchestrator_chain = ChainNumCollators {
49
1
        para_id: 1000.into(),
50
1
        min_collators: 2,
51
1
        max_collators: 5,
52
1
        parathread: false,
53
1
    };
54
1
    let mut old_assigned = BTreeMap::from_iter(vec![(1000.into(), vec![1, 2])]);
55
1

            
56
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
57
1
        &collators,
58
1
        orchestrator_chain,
59
1
        &mut old_assigned,
60
1
    );
61
1

            
62
1
    assert_eq!(num_invulnerables, 0);
63
1
}
64

            
65
#[test]
66
1
fn invulnerable_priority_1_invulnerable_orchestrator() {
67
1
    let collators = vec![1, 2, 3, 4, 5, 101];
68
1
    let orchestrator_chain = ChainNumCollators {
69
1
        para_id: 1000.into(),
70
1
        min_collators: 2,
71
1
        max_collators: 5,
72
1
        parathread: false,
73
1
    };
74
1
    let mut old_assigned = BTreeMap::from_iter(vec![(1000.into(), vec![101])]);
75
1

            
76
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
77
1
        &collators,
78
1
        orchestrator_chain,
79
1
        &mut old_assigned,
80
1
    );
81
1

            
82
1
    assert_eq!(num_invulnerables, 1);
83
1
}
84

            
85
#[test]
86
1
fn invulnerable_priority_1_invulnerable_not_assigned() {
87
1
    let collators = vec![1, 2, 3, 4, 5, 101];
88
1
    let orchestrator_chain = ChainNumCollators {
89
1
        para_id: 1000.into(),
90
1
        min_collators: 2,
91
1
        max_collators: 5,
92
1
        parathread: false,
93
1
    };
94
1
    let mut old_assigned = BTreeMap::new();
95
1

            
96
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
97
1
        &collators,
98
1
        orchestrator_chain,
99
1
        &mut old_assigned,
100
1
    );
101
1

            
102
1
    assert_eq!(num_invulnerables, 1);
103
1
}
104

            
105
#[test]
106
1
fn invulnerable_priority_1_invulnerable_assigned_to_another_chain() {
107
1
    let collators = vec![1, 2, 3, 4, 5, 101];
108
1
    let orchestrator_chain = ChainNumCollators {
109
1
        para_id: 1000.into(),
110
1
        min_collators: 2,
111
1
        max_collators: 5,
112
1
        parathread: false,
113
1
    };
114
1
    let mut old_assigned =
115
1
        BTreeMap::from_iter(vec![(1000.into(), vec![]), (2000.into(), vec![101])]);
116
1

            
117
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
118
1
        &collators,
119
1
        orchestrator_chain,
120
1
        &mut old_assigned,
121
1
    );
122
1

            
123
1
    assert_eq!(num_invulnerables, 1);
124
1
}
125

            
126
#[test]
127
1
fn bug_same_invulnerable_selected_twice() {
128
1
    let collators = vec![100];
129
1
    let orchestrator_chain = ChainNumCollators {
130
1
        para_id: 1000.into(),
131
1
        min_collators: 2,
132
1
        max_collators: 5,
133
1
        parathread: false,
134
1
    };
135
1
    let mut old_assigned = BTreeMap::from_iter(vec![(1000.into(), vec![100])]);
136
1

            
137
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
138
1
        &collators,
139
1
        orchestrator_chain,
140
1
        &mut old_assigned,
141
1
    );
142
1

            
143
1
    assert_eq!(num_invulnerables, 1);
144
1
}
145

            
146
#[test]
147
1
fn bug_not_using_assigned_invulnerables() {
148
1
    // There are 3 invulnerables, 1 assigned to orchestrator and 2 assigned to a container chain.
149
1
    // After `prioritize_invulnerables` the first one from the container should move to orchestrator
150
1
    let collators = vec![1, 2, 3, 4, 5, 102, 103, 104];
151
1

            
152
1
    let container_chains = [
153
1
        ChainNumCollators {
154
1
            para_id: 1000.into(),
155
1
            min_collators: 2,
156
1
            max_collators: 5,
157
1
            parathread: false,
158
1
        },
159
1
        ChainNumCollators {
160
1
            para_id: 2000.into(),
161
1
            min_collators: 2,
162
1
            max_collators: 2,
163
1
            parathread: false,
164
1
        },
165
1
        ChainNumCollators {
166
1
            para_id: 2001.into(),
167
1
            min_collators: 2,
168
1
            max_collators: 2,
169
1
            parathread: false,
170
1
        },
171
1
    ];
172
1
    let orchestrator_chain = container_chains[0];
173
1

            
174
1
    let mut old_assigned = BTreeMap::from_iter(vec![
175
1
        (1000.into(), vec![101, 104, 1, 2, 3]),
176
1
        (2000.into(), vec![4, 5]),
177
1
        (2001.into(), vec![102, 103]),
178
1
    ]);
179
1

            
180
1
    let chains_with_collators_set =
181
3
        BTreeSet::from_iter(container_chains.iter().map(|cc| cc.para_id));
182
1
    let collators_set = BTreeSet::from_iter(collators.iter().cloned());
183
1
    Assignment::<Test>::retain_valid_old_assigned(
184
1
        &mut old_assigned,
185
1
        &chains_with_collators_set,
186
1
        &collators_set,
187
1
    );
188
1
    let num_invulnerables = Assignment::<Test>::prioritize_invulnerables(
189
1
        &collators,
190
1
        orchestrator_chain,
191
1
        &mut old_assigned,
192
1
    );
193
1

            
194
1
    assert_eq!(num_invulnerables, 2);
195
1
}