More tests and a small cleanup

dfinity · Oct 18, 2024 · 7f1dd0d · 7f1dd0d
1 parent 673ce0e
commit 7f1dd0d
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 1 deletion.
diff --git a/rs/tla_instrumentation/tla_instrumentation/src/tla_state.rs b/rs/tla_instrumentation/tla_instrumentation/src/tla_state.rs
@@ -1,6 +1,5 @@
 use crate::tla_value::{TlaValue, ToTla};
 use candid::CandidType;
-// use serde::Deserialize;
 use crate::SourceLocation;
 use std::{
     collections::{BTreeMap, BTreeSet},

diff --git a/rs/tla_instrumentation/tla_instrumentation/tests/multiple_calls.rs b/rs/tla_instrumentation/tla_instrumentation/tests/multiple_calls.rs
@@ -280,5 +280,49 @@ fn multiple_calls_test() {
         )
     );
 
+    let third = &pairs[2];
+
+    assert_eq!(third.start.get("counter"), Some(&2_u64.to_tla_value()));
+    assert_eq!(third.end.get("counter"), Some(&3_u64.to_tla_value()));
+
+    assert_eq!(
+        third.start.get("my_local"),
+        Some(BTreeMap::from([(PID, 2_u64)]).to_tla_value()).as_ref()
+    );
+    assert_eq!(
+        third.end.get("my_local"),
+        Some(BTreeMap::from([(PID, 3_u64)]).to_tla_value()).as_ref()
+    );
+
+    assert_eq!(
+        third.start.get(incoming),
+        Some(
+            &BTreeSet::from([TlaValue::Record(BTreeMap::from([
+                ("caller".to_string(), PID.to_tla_value()),
+                (
+                    "response".to_string(),
+                    TlaValue::Variant {
+                        tag: "Ok".to_string(),
+                        value: Box::new(3_u64.to_tla_value())
+                    }
+                )
+            ]))])
+            .to_tla_value()
+        )
+    );
+    assert_eq!(
+        third.end.get(incoming),
+        Some(&BTreeSet::<TlaValue>::new().to_tla_value())
+    );
+
+    assert_eq!(
+        third.start.get(outgoing),
+        Some(&Vec::<TlaValue>::new().to_tla_value())
+    );
+    assert_eq!(
+        third.end.get(outgoing),
+        Some(&Vec::<TlaValue>::new().to_tla_value())
+    );
+
     check_tla_trace(trace);
 }