Add test case for the Host version of get_qubit_by_id

Author: Stephen DiAdamo

integration_tests/test_host.py

@@ -1,5 +1,7 @@
 import unittest
-from qunetsim.components.host import Host
+
+from qunetsim.objects import Qubit
+from qunetsim.components import Host
 from random import randint
 
 
@@ -15,6 +17,32 @@ def setUpClass(cls):
     def tearDownClass(cls):
         pass
 
+    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'))
+
     # @unittest.skip('')
     def test_connections(self):
         name = "A"