Improve token-specified login flow. (#169)

* Test that we can create multiple IBMQ devices using an environment token. * Improve token-specified login flow. * Fix formatting. * extend skipping logic * Test token login elision. * Temporarily restrict code coverage analysis. * Update changelog. * Fix formatting. * more no cover for the IBMQ device * Update tests/test_ibmq.py Co-authored-by: antalszava <[email protected]> * Update tests/test_ibmq.py Co-authored-by: antalszava <[email protected]> * Update test_ibmq.py Co-authored-by: Antal Szava <[email protected]>
PennyLaneAI · Dec 15, 2021 · ccfdf34 · ccfdf34
1 parent d5eea9b
commit ccfdf34
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,9 @@
 
 ### Improvements
 
+* Improved the login flow when IBMQ tokens are specified as environement variables.
+  [(#169)](https://github.com/PennyLaneAI/pennylane-qiskit/pull/169)
+
 ### Documentation
 
 ### Bug fixes
@@ -15,6 +18,8 @@
 
 This release contains contributions from (in alphabetical order):
 
+Tanner Rogalsky
+
 ---
 
 # Release 0.20.0

diff --git a/pennylane_qiskit/ibmq.py b/pennylane_qiskit/ibmq.py
@@ -66,11 +66,24 @@ def __init__(self, wires, provider=None, backend="ibmq_qasm_simulator", shots=10
         group = kwargs.get("group", "open")
         project = kwargs.get("project", "main")
 
-        if token is not None:
+        # TODO: remove "no cover" when #173 is resolved
+        if token is not None:  # pragma: no cover
             # token was provided by the user, so attempt to enable an
             # IBM Q account manually
-            ibmq_kwargs = {"url": url} if url is not None else {}
-            IBMQ.enable_account(token, **ibmq_kwargs)
+            def login():
+                ibmq_kwargs = {"url": url} if url is not None else {}
+                IBMQ.enable_account(token, **ibmq_kwargs)
+
+            active_account = IBMQ.active_account()
+            if active_account is None:
+                login()
+            else:
+                # There is already an active account:
+                # If the token is the same, do nothing.
+                # If the token is different, authenticate with the new account.
+                if active_account["token"] != token:
+                    IBMQ.disable_account()
+                    login()
         else:
             # check if an IBM Q account is already active.
             #
@@ -96,13 +109,13 @@ def __init__(self, wires, provider=None, backend="ibmq_qasm_simulator", shots=10
 
         super().__init__(wires=wires, provider=p, backend=backend, shots=shots, **kwargs)
 
-    def batch_execute(self, circuits):
+    def batch_execute(self, circuits):  # pragma: no cover
         res = super().batch_execute(circuits)
         if self.tracker.active:
             self._track_run()
         return res
 
-    def _track_run(self):
+    def _track_run(self):  # pragma: no cover
         """Provide runtime information."""
 
         time_per_step = self._current_job.time_per_step()

diff --git a/tests/test_ibmq.py b/tests/test_ibmq.py
@@ -18,7 +18,7 @@ def token():
     variable."""
     t = os.getenv("IBMQX_TOKEN_TEST", None)
 
-    if t is None:
+    if t is None or t == '':
         pytest.skip("Skipping test, no IBMQ token available")
 
     yield t
@@ -32,6 +32,47 @@ def test_load_from_env(token, monkeypatch):
     dev = IBMQDevice(wires=1)
     assert dev.provider.credentials.is_ibmq()
 
+def test_load_from_env_multiple_device(token, monkeypatch):
+    """Test creating multiple IBMQ devices when the environment variable
+    for the IBMQ token was set."""
+    monkeypatch.setenv("IBMQX_TOKEN", token)
+    dev1 = IBMQDevice(wires=1)
+    dev2 = IBMQDevice(wires=1)
+
+    assert dev1.provider.credentials.is_ibmq()
+    assert dev2.provider.credentials.is_ibmq()
+
+def test_load_from_env_multiple_device_and_token(monkeypatch):
+    """Test creating multiple devices when the different tokens are defined
+    using an environment variable."""
+    mock_provider = "MockProvider"
+    mock_qiskit_device = MockQiskitDeviceInit()
+
+    with monkeypatch.context() as m:
+        m.setattr(ibmq.QiskitDevice, "__init__", mock_qiskit_device.mocked_init)
+
+        creds = []
+        def enable_account(new_creds):
+            creds.append(new_creds)
+        def active_account():
+            if len(creds) != 0:
+                return { "token": creds[-1] }
+        m.setattr(ibmq.IBMQ, "enable_account", enable_account)
+        m.setattr(ibmq.IBMQ, "disable_account", lambda: None)
+        m.setattr(ibmq.IBMQ, "active_account", active_account)
+
+        m.setenv("IBMQX_TOKEN", "TOKEN1")
+        dev1 = IBMQDevice(wires=1, provider=mock_provider)
+        # first login
+        assert creds == ["TOKEN1"]
+        dev1 = IBMQDevice(wires=1, provider=mock_provider)
+        # same token, login is elided
+        assert creds == ["TOKEN1"]
+
+        m.setenv("IBMQX_TOKEN", "TOKEN2")
+        dev2 = IBMQDevice(wires=1, provider=mock_provider)
+        # second login
+        assert creds == ["TOKEN1", "TOKEN2"]
 
 def test_load_kwargs_takes_precedence(token, monkeypatch):
     """Test that with a potentially valid token stored as an environment