From 1d4d5f4c6b11d9e5335cf3c2cab3feec2bffa3af Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <107005333+laurencap@users.noreply.github.com>
Date: Wed, 21 Jun 2023 14:51:03 -0400
Subject: [PATCH] [autoqasm] Increase qubits.py coverage to 100% (#283)

---
 .../experimental/autoqasm/gates/qubits.py     | 29 +++++++++-------
 .../braket/experimental/autoqasm/test_api.py  | 33 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/src/braket/experimental/autoqasm/gates/qubits.py b/src/braket/experimental/autoqasm/gates/qubits.py
index 7bb11a06d..416d0f6a8 100644
--- a/src/braket/experimental/autoqasm/gates/qubits.py
+++ b/src/braket/experimental/autoqasm/gates/qubits.py
@@ -45,6 +45,23 @@ def _qubit(qid: Any) -> oqpy.Qubit:
     raise ValueError(f"invalid qubit label: '{qid}'")
 
 
+@_qubit.register
+def _(qid: bool) -> oqpy.Qubit:
+    raise ValueError(f"invalid qubit label: '{qid}'")
+
+
+@_qubit.register
+def _(qid: float) -> oqpy.Qubit:
+    raise TypeError(f"qubit index cannot be a float: '{qid}'")
+
+
+@_qubit.register
+def _(qid: int) -> oqpy.Qubit:
+    # Integer virtual qubit, like `h(0)`
+    program.get_program_conversion_context().register_qubit(qid)
+    return oqpy.Qubit(_global_qubit_register(qid), needs_declaration=False)
+
+
 @_qubit.register
 def _(qid: oqpy._ClassicalVar) -> oqpy.Qubit:
     # Indexed by variable, such as i in range(n); h(i)
@@ -64,13 +81,6 @@ def _(qid: oqpy.base.OQPyExpression) -> oqpy.Qubit:
     return oqpy.Qubit(_global_qubit_register(qubit_idx_expr), needs_declaration=False)
 
 
-@_qubit.register
-def _(qid: int) -> oqpy.Qubit:
-    # Integer virtual qubit, like `h(0)`
-    program.get_program_conversion_context().register_qubit(qid)
-    return oqpy.Qubit(_global_qubit_register(qid), needs_declaration=False)
-
-
 @_qubit.register
 def _(qid: str) -> oqpy.Qubit:
     # Physical qubit label, like `h("$0")`
@@ -83,8 +93,3 @@ def _(qid: str) -> oqpy.Qubit:
         return oqpy.PhysicalQubits[qubit_idx]
     else:
         raise ValueError(f"invalid qubit label: '{qid}'")
-
-
-@_qubit.register
-def _(qid: float) -> oqpy.Qubit:
-    raise TypeError(f"qubit index cannot be a float: '{qid}'")
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_api.py b/test/unit_tests/braket/experimental/autoqasm/test_api.py
index c55597d81..91fed86d9 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_api.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_api.py
@@ -560,6 +560,30 @@ def broken() -> None:
         broken()
 
 
+def test_bool_qubit_index_fails() -> None:
+    """Tests that an error is raised for boolean qubit type."""
+
+    @aq.function
+    def broken() -> None:
+        "Uses invalid type for qubit index"
+        h(True)
+
+    with pytest.raises(ValueError):
+        broken()
+
+
+def test_invalid_qubit_type_fails() -> None:
+    """Tests that an error is raised for other unusual qubit types."""
+
+    @aq.function
+    def broken() -> None:
+        "Uses invalid type for qubit index"
+        h(h)
+
+    with pytest.raises(ValueError):
+        broken()
+
+
 def test_bit_array_name() -> None:
     """Tests that auto declared bits are given a reasonable name."""
 
@@ -632,8 +656,15 @@ def test_program_with_expr() -> None:
     """Test that a program with expressions for the qubit index raises
     an error if the user doesn't specify the number of qubits.
     """
+
+    @aq.function
+    def qubit_expr() -> None:
+        "Uses aq.range iterator for qubit index."
+        for i in aq.range(5):
+            h(i + 3)
+
     with pytest.raises(errors.UnknownQubitCountError):
-        ghz_qasm_for_loop()
+        qubit_expr()
 
 
 def test_program_with_expr_and_declaration() -> None: