def test_get_qubits_by_id(self):
host = Host('A')
q1 = Qubit(host)
q2 = Qubit(host)
q3 = Qubit(host)
host2 = Host('B')
q4 = Qubit(host2)
host.add_epr('B', q1)
host.add_data_qubit('C', q2)
host.add_ghz_qubit('D', q3)
host2.add_data_qubit('A', q4)
# Test all types of qubits
self.assertEqual(q1, host.get_qubit_by_id(q1.id))
self.assertEqual(q2, host.get_qubit_by_id(q2.id))
self.assertEqual(q3, host.get_qubit_by_id(q3.id))
# Test getting qubits from other hosts
self.assertIsNone(host.get_qubit_by_id(q4.id))
# Test getting qubits that don't exist
self.assertIsNone(host.get_qubit_by_id('fake'))
def test_resetting_qubits(self):
host = Host('A')
q1 = Qubit(host)
q2 = Qubit(host)
q3 = Qubit(host)
host.add_epr('B', q1)
host.add_data_qubit('B', q2)
host.add_ghz_qubit('B', q3)
qs = host.get_data_qubits('B')
self.assertEqual(len(qs), 1)
host.reset_data_qubits('B')
qs = host.get_data_qubits('B')
self.assertEqual(len(qs), 0)
