From 91f690d49526f9b96840e9eb97499f73f92b68a5 Mon Sep 17 00:00:00 2001
From: Spencer Churchill <churchill@ionq.co>
Date: Wed, 26 Jul 2023 15:08:48 -0700
Subject: [PATCH] Add cirq specific environment variables to ionq (#6216)

* add cirq specific environment variables

* use mock.patch.dict to set env var

* add tests for env vars

* add tests for remote_host + use decorators

---------

Co-authored-by: Tanuj Khattar <tanujkhattar@google.com>
---
 cirq-ionq/cirq_ionq/service.py      |  7 ++-
 cirq-ionq/cirq_ionq/service_test.py | 66 ++++++++++++++++++++++++++---
 2 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/cirq-ionq/cirq_ionq/service.py b/cirq-ionq/cirq_ionq/service.py
index 1a052e2d96d..1b2384e2177 100644
--- a/cirq-ionq/cirq_ionq/service.py
+++ b/cirq-ionq/cirq_ionq/service.py
@@ -62,9 +62,12 @@ def __init__(
                 This is actually an EnvironmentError which is equal to an OSError.
         """
         self.remote_host = (
-            remote_host or os.getenv('IONQ_REMOTE_HOST') or f'https://api.ionq.co/{api_version}'
+            remote_host
+            or os.getenv('CIRQ_IONQ_REMOTE_HOST')
+            or os.getenv('IONQ_REMOTE_HOST')
+            or f'https://api.ionq.co/{api_version}'
         )
-        self.api_key = api_key or os.getenv('IONQ_API_KEY')
+        self.api_key = api_key or os.getenv('CIRQ_IONQ_API_KEY') or os.getenv('IONQ_API_KEY')
         if not self.api_key:
             raise EnvironmentError(
                 'Parameter api_key was not specified and the environment variable '
diff --git a/cirq-ionq/cirq_ionq/service_test.py b/cirq-ionq/cirq_ionq/service_test.py
index 3d4d62a4691..c673d562cab 100644
--- a/cirq-ionq/cirq_ionq/service_test.py
+++ b/cirq-ionq/cirq_ionq/service_test.py
@@ -158,25 +158,81 @@ def test_service_list_calibrations():
     mock_client.list_calibrations.assert_called_with(start=start, end=end, limit=10, batch_size=2)
 
 
+@mock.patch.dict(os.environ, {'IONQ_API_KEY': 'tomyheart'})
 def test_service_api_key_via_env():
-    os.environ['IONQ_API_KEY'] = 'tomyheart'
     service = ionq.Service(remote_host='http://example.com')
     assert service.api_key == 'tomyheart'
-    del os.environ['IONQ_API_KEY']
 
 
+@mock.patch.dict(os.environ, {'IONQ_REMOTE_HOST': 'http://example.com'})
 def test_service_remote_host_via_env():
-    os.environ['IONQ_REMOTE_HOST'] = 'http://example.com'
     service = ionq.Service(api_key='tomyheart')
     assert service.remote_host == 'http://example.com'
-    del os.environ['IONQ_REMOTE_HOST']
 
 
+@mock.patch.dict(os.environ, {}, clear=True)
 def test_service_no_param_or_env_variable():
     with pytest.raises(EnvironmentError):
         _ = ionq.Service(remote_host='http://example.com')
 
 
-def test_service_no_url_default():
+@mock.patch.dict(os.environ, {'IONQ_API_KEY': 'not_this_key'})
+def test_service_api_key_passed_directly():
+    service = ionq.Service(remote_host='http://example.com', api_key='tomyheart')
+    assert service.api_key == 'tomyheart'
+
+
+@mock.patch.dict(os.environ, {'CIRQ_IONQ_API_KEY': 'tomyheart'})
+def test_service_api_key_from_env_var_cirq_ionq():
+    service = ionq.Service(remote_host='http://example.com')
+    assert service.api_key == 'tomyheart'
+
+
+@mock.patch.dict(os.environ, {'IONQ_API_KEY': 'tomyheart'})
+def test_service_api_key_from_env_var_ionq():
+    service = ionq.Service(remote_host='http://example.com')
+    assert service.api_key == 'tomyheart'
+
+
+@mock.patch.dict(os.environ, {}, clear=True)
+def test_service_api_key_not_found_raises_error():
+    with pytest.raises(EnvironmentError):
+        _ = ionq.Service(remote_host='http://example.com')
+
+
+@mock.patch.dict(os.environ, {'CIRQ_IONQ_API_KEY': 'tomyheart', 'IONQ_API_KEY': 'not_this_key'})
+def test_service_api_key_from_env_var_cirq_ionq_precedence():
+    service = ionq.Service(remote_host='http://example.com')
+    assert service.api_key == 'tomyheart'
+
+
+@mock.patch.dict(os.environ, {'CIRQ_IONQ_REMOTE_HOST': 'not_this_host'})
+def test_service_remote_host_passed_directly():
+    service = ionq.Service(remote_host='http://example.com', api_key='tomyheart')
+    assert service.remote_host == 'http://example.com'
+
+
+@mock.patch.dict(os.environ, {'CIRQ_IONQ_REMOTE_HOST': 'http://example.com'})
+def test_service_remote_host_from_env_var_cirq_ionq():
+    service = ionq.Service(api_key='tomyheart')
+    assert service.remote_host == 'http://example.com'
+
+
+@mock.patch.dict(os.environ, {'IONQ_REMOTE_HOST': 'http://example.com'})
+def test_service_remote_host_from_env_var_ionq():
     service = ionq.Service(api_key='tomyheart')
+    assert service.remote_host == 'http://example.com'
+
+
+@mock.patch.dict(os.environ, {}, clear=True)
+def test_service_remote_host_default():
+    service = ionq.Service(api_key='tomyheart', api_version='v0.1')
     assert service.remote_host == 'https://api.ionq.co/v0.1'
+
+
+@mock.patch.dict(
+    os.environ, {'CIRQ_IONQ_REMOTE_HOST': 'http://example.com', 'IONQ_REMOTE_HOST': 'not_this_host'}
+)
+def test_service_remote_host_from_env_var_cirq_ionq_precedence():
+    service = ionq.Service(api_key='tomyheart')
+    assert service.remote_host == 'http://example.com'