From 39ba5d4dea1176b66a3b6421032726059808584b Mon Sep 17 00:00:00 2001
From: Andrea Papaluca <andreapapaluca@gmail.com>
Date: Mon, 11 Mar 2024 13:15:47 +0400
Subject: [PATCH] feat: added test for moments with callback gates

---
 src/qibo/models/circuit.py   | 2 +-
 tests/test_models_circuit.py | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index b06f118e7c..def7a7f6d4 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -94,7 +94,7 @@ def moments(self):
         moments = [self.nqubits * [None]]
         moment_index = self.nqubits * [0]
         for gate in self:
-            if gate.qubits:
+            if not isinstance(gate, gates.CallbackGate):
                 qubits = gate.qubits
             else:  # special gate acting on all qubits
                 qubits = tuple(range(self.nqubits))
diff --git a/tests/test_models_circuit.py b/tests/test_models_circuit.py
index 8d83b9d83e..5181f57a61 100644
--- a/tests/test_models_circuit.py
+++ b/tests/test_models_circuit.py
@@ -19,8 +19,10 @@ def test_parametrizedgates_class():
 
 
 def test_queue_class():
+    from qibo.callbacks import EntanglementEntropy
     from qibo.models.circuit import _Queue
 
+    entropy = EntanglementEntropy([0])
     queue = _Queue(4)
     gatelist = [
         gates.H(0),
@@ -29,12 +31,14 @@ def test_queue_class():
         gates.H(2),
         gates.CNOT(1, 2),
         gates.Y(3),
+        gates.CallbackGate(entropy),
     ]
     for g in gatelist:
         queue.append(g)
     assert queue.moments == [
         [gatelist[0], gatelist[1], gatelist[3], gatelist[5]],
         [gatelist[2], gatelist[4], gatelist[4], None],
+        [gatelist[6] for _ in range(4)],
     ]