From 273ce34958ca73c865b1e317074899127404595b Mon Sep 17 00:00:00 2001
From: Andy Leiserson <aleiserson@mozilla.com>
Date: Wed, 9 Nov 2022 15:01:51 -0800
Subject: [PATCH] tests for #135

---
 src/protocol/context.rs | 68 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/src/protocol/context.rs b/src/protocol/context.rs
index d2b744f8e..577b78f4e 100644
--- a/src/protocol/context.rs
+++ b/src/protocol/context.rs
@@ -193,3 +193,71 @@ impl<'a, F: Field> ProtocolContext<'a, MaliciousReplicated<F>, F> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        ff::Fp31,
+        protocol::QueryId,
+        test_fixture::{make_contexts, make_world},
+    };
+
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    struct TestStep;
+
+    impl Substep for TestStep {}
+
+    impl AsRef<str> for TestStep {
+        fn as_ref(&self) -> &str {
+            "test"
+        }
+    }
+
+    #[tokio::test]
+    pub async fn narrow() {
+        let world = make_world(QueryId);
+        let context = make_contexts::<Fp31>(&world);
+
+        let _ctx = context[0].narrow(&TestStep);
+    }
+
+    // In debug builds, we should panic if we try to reuse a step in the same context.
+    #[cfg(debug_assertions)]
+    #[tokio::test]
+    #[should_panic]
+    pub async fn narrow_twice() {
+        let world = make_world(QueryId);
+        let context = make_contexts::<Fp31>(&world);
+
+        let _ctx = context[0].narrow(&TestStep);
+        let _ctx = context[0].narrow(&TestStep);
+    }
+
+    // When a context is narrowed inside a loop, that should not be reported as a
+    // duplicate usage. Note, however, that care must be taken when writing protocols
+    // or this exception may let things through that are not okay. For example, the
+    // following is okay:
+    //
+    // ```
+    // for i in 0..2 {
+    //  // do something for RecordId i
+    // }
+    // ```
+    //
+    // ```
+    // for i in 0..2 {
+    //  // do something *without* narrowing context by `i` or using it as RecordId
+    // }
+    // ```
+    #[ignore] // broken, see #135
+    #[tokio::test]
+    pub async fn narrow_in_loop() {
+        let world = make_world(QueryId);
+        let context = make_contexts::<Fp31>(&world);
+
+        for _ in 0..2 {
+            let _ctx = context[0].narrow(&TestStep);
+        }
+    }
+}