1use std::collections::BTreeSet;
2
3use trellis_core::{ResourceCommandTrace, ResourceKey, Revision, ScopeId};
4
5use crate::{
6 HostStatusClass, HostStatusRecord, ResourceLedger, ResourceLedgerError, ResourceSnapshot,
7 ResourceStatusContext,
8};
9
10impl<C> ResourceLedger<C> {
11 pub fn assert_command_order(
13 &self,
14 expected: &[ResourceCommandTrace],
15 ) -> Result<(), ResourceLedgerError> {
16 if self.command_trace == expected {
17 Ok(())
18 } else {
19 Err(ResourceLedgerError::CommandOrderMismatch {
20 expected: expected.to_vec(),
21 actual: self.command_trace.clone(),
22 })
23 }
24 }
25
26 pub fn assert_all_resources_have_owner(&self) -> Result<(), ResourceLedgerError> {
28 for (key, snapshot) in &self.resources {
29 if snapshot.owners.is_empty() {
30 return Err(ResourceLedgerError::Orphan {
31 key: key.clone(),
32 context: Some(snapshot.command_context()),
33 });
34 }
35 }
36 Ok(())
37 }
38
39 pub fn assert_no_orphan_resources(&self) -> Result<(), ResourceLedgerError> {
41 self.assert_all_resources_have_owner()
42 }
43
44 pub fn assert_no_duplicate_close(&self) -> Result<(), ResourceLedgerError> {
46 if let Some(context) = self.duplicate_closes.first() {
47 Err(ResourceLedgerError::DuplicateClose {
48 key: context.key.clone(),
49 context: context.clone(),
50 })
51 } else {
52 Ok(())
53 }
54 }
55
56 pub fn assert_no_forbidden_opened(&self) -> Result<(), ResourceLedgerError> {
58 if let Some(context) = self.forbidden_opened.first() {
59 Err(ResourceLedgerError::ForbiddenOpen {
60 key: context.key.clone(),
61 context: Some(context.clone()),
62 })
63 } else {
64 Ok(())
65 }
66 }
67
68 pub fn assert_no_wildcard_resource_opened(&self) -> Result<(), ResourceLedgerError> {
70 self.assert_no_forbidden_opened()
71 }
72
73 pub fn assert_resource_not_open(&self, key: &ResourceKey) -> Result<(), ResourceLedgerError> {
75 if self.resources.contains_key(key) {
76 Err(ResourceLedgerError::StillOpen {
77 key: key.clone(),
78 context: self.context_for_key(key),
79 })
80 } else {
81 Ok(())
82 }
83 }
84
85 pub fn assert_closed_scope_owns_no_resources(
87 &self,
88 scope: ScopeId,
89 ) -> Result<(), ResourceLedgerError> {
90 let resources = self
91 .resources
92 .iter()
93 .filter(|(_, snapshot)| snapshot.owners.contains(&scope))
94 .map(|(key, _)| key.clone())
95 .collect::<Vec<_>>();
96 if resources.is_empty() {
97 Ok(())
98 } else {
99 let contexts = resources
100 .iter()
101 .filter_map(|key| self.context_for_key(key))
102 .collect();
103 Err(ResourceLedgerError::ClosedScopeOwnsResources {
104 scope,
105 resources,
106 contexts,
107 })
108 }
109 }
110
111 pub fn assert_resource_opened_once(
113 &self,
114 key: &ResourceKey,
115 ) -> Result<(), ResourceLedgerError> {
116 self.assert_count(key, "open_count", 1, |snapshot| snapshot.open_count)
117 }
118
119 pub fn assert_resource_closed_once(
121 &self,
122 key: &ResourceKey,
123 ) -> Result<(), ResourceLedgerError> {
124 self.assert_count(key, "close_count", 1, |snapshot| snapshot.close_count)
125 }
126
127 pub fn assert_resource_generation(
129 &self,
130 key: &ResourceKey,
131 expected: u64,
132 ) -> Result<(), ResourceLedgerError> {
133 let actual = self
134 .history
135 .get(key)
136 .map_or(0, |snapshot| snapshot.generation);
137 if actual == expected {
138 Ok(())
139 } else {
140 Err(ResourceLedgerError::GenerationMismatch {
141 key: key.clone(),
142 expected,
143 actual,
144 context: self.context_for_key(key),
145 })
146 }
147 }
148
149 pub fn assert_resource_shared_by(
151 &self,
152 key: &ResourceKey,
153 expected: BTreeSet<ScopeId>,
154 ) -> Result<(), ResourceLedgerError> {
155 let actual = self
156 .resources
157 .get(key)
158 .map(|snapshot| snapshot.owners.clone())
159 .unwrap_or_default();
160 if actual == expected {
161 Ok(())
162 } else {
163 Err(ResourceLedgerError::OwnerMismatch {
164 key: key.clone(),
165 expected,
166 actual,
167 context: self.context_for_key(key),
168 })
169 }
170 }
171
172 pub fn assert_status_is_stale(
174 &self,
175 key: &ResourceKey,
176 command_revision: Revision,
177 ) -> Result<(), ResourceLedgerError> {
178 let Some(record) = self.status_records.iter().find(|record| {
179 record.status.resource_key == *key && record.status.command_revision == command_revision
180 }) else {
181 return Err(ResourceLedgerError::MissingStatus {
182 key: key.clone(),
183 command_revision,
184 });
185 };
186 if record.class == HostStatusClass::Stale {
187 Ok(())
188 } else {
189 Err(ResourceLedgerError::StatusClassMismatch {
190 context: status_context(record),
191 expected: HostStatusClass::Stale,
192 })
193 }
194 }
195
196 pub fn assert_status_did_not_resurrect_closed_scope(
198 &self,
199 scope: ScopeId,
200 ) -> Result<(), ResourceLedgerError> {
201 self.assert_closed_scope_owns_no_resources(scope)?;
202 self.assert_no_status_mutated_closed_scope()
203 }
204
205 pub fn assert_no_status_mutated_closed_scope(&self) -> Result<(), ResourceLedgerError> {
207 for record in &self.status_records {
208 if record.class == HostStatusClass::Late
209 && self.scope_owns_resource(record.status.scope, &record.status.resource_key)
210 {
211 return Err(ResourceLedgerError::StatusMutatedClosedScope {
212 scope: record.status.scope,
213 context: status_context(record),
214 });
215 }
216 }
217 Ok(())
218 }
219
220 fn assert_count(
221 &self,
222 key: &ResourceKey,
223 field: &'static str,
224 expected: usize,
225 count: impl FnOnce(&ResourceSnapshot<C>) -> usize,
226 ) -> Result<(), ResourceLedgerError> {
227 let actual = self.history.get(key).map_or(0, count);
228 if actual == expected {
229 Ok(())
230 } else {
231 Err(ResourceLedgerError::CountMismatch {
232 key: key.clone(),
233 field,
234 expected,
235 actual,
236 context: self.context_for_key(key),
237 })
238 }
239 }
240
241 fn scope_owns_resource(&self, scope: ScopeId, key: &ResourceKey) -> bool {
242 self.resources
243 .get(key)
244 .is_some_and(|snapshot| snapshot.owners.contains(&scope))
245 }
246}
247
248fn status_context(record: &HostStatusRecord) -> ResourceStatusContext {
249 ResourceStatusContext {
250 status: record.status.clone(),
251 class: record.class,
252 last_transaction_id: record.last_transaction_id,
253 last_command_revision: record.last_command_revision,
254 }
255}