Add test case for passing a callable into an aggregate exp fragment

OxfordIonTrapGroup · Dec 9, 2024 · ce2268c · ce2268c
1 parent df0d018
commit ce2268c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/test/fixtures.py b/test/fixtures.py
@@ -202,7 +202,13 @@ class AddOneAggregate(AggregateExpFragment):
     def build_fragment(self) -> None:
         self.setattr_fragment("a", AddOneFragment)
         self.setattr_fragment("b", AddOneFragment)
-        return super().build_fragment([self.a, self.b])
+
+        self.setattr_result("sum", FloatChannel)
+
+        def push_sum():
+            self.sum.push(self.a.value.get() + 1 + self.b.value.get() + 1)
+
+        return super().build_fragment([self.a, self.b, push_sum])
 
 
 class TrivialKernelAggregate(AggregateExpFragment):

diff --git a/test/test_experiment_entrypoint.py b/test/test_experiment_entrypoint.py
@@ -32,6 +32,7 @@ def test_aggregate(self):
         self.assertEqual(parent.b.num_prepare_calls, 1)
         self.assertEqual(result[parent.a.result], 1.0)
         self.assertEqual(result[parent.b.result], 1.0)
+        self.assertEqual(result[parent.sum], 2.0)
 
     def test_kernel(self):
         parent = self.create(TrivialKernelAggregate, [])