Expose authenticate_user

cirq-google/cirq_google/engine/qcs_notebook.py

@@ -82,32 +82,14 @@ def get_qcs_objects_for_notebook(
         ValueError: if processor_id is not specified and no processors are available.
     """
 
-    # Check for Google Application Default Credentials and run
-    # interactive login if the notebook is executed in Colab. In
-    # case the notebook is executed in Jupyter notebook or other
-    # IPython runtimes, no interactive login is provided, it is
-    # assumed that the `GOOGLE_APPLICATION_CREDENTIALS` env var is
-    # set or `gcloud auth application-default login` was executed
-    # already. For more information on using Application Default Credentials
-    # see https://cloud.google.com/docs/authentication/production
-    # Attempt to connect to the Quantum Engine API, and use a simulator if unable to connect.
     if not virtual:
         # Set up auth
         try:
-            from google.colab import auth
-        except ImportError:
-            print("Not running in a colab kernel. Will use Application Default Credentials.")
-        else:
-            print("Getting OAuth2 credentials.")
-            print("Press enter after entering the verification code.")
-            try:
-                a = auth.authenticate_user(clear_output=False)
-                print(a)
-                print("Authentication complete.")
-            except Exception as exc:
-                print(f"Authentication failed: {exc}")
-                print("Using virtual engine instead.")
-                virtual = True
+            authenticate_user()
+        except Exception as exc:
+            print(f"Authentication failed: {exc}")
+            print("Using virtual engine instead.")
+            virtual = True
 
     if not virtual:
         # Set up production engine
@@ -152,3 +134,46 @@ def get_qcs_objects_for_notebook(
         processor_id=processor_id,
         is_simulator=is_simulator,
     )
+
+
+def authenticate_user(clear_output: bool = False) -> None:
+    """Authenticates on Google Cloud.
+
+    Args:
+        clear_output: Optional bool for whether to clear output before
+            authenticating. Defaults to false.
+
+    Returns:
+        None.
+
+    Raises:
+        Exception: if authentication fails.
+    """
+
+    # Check for Google Application Default Credentials and run
+    # interactive login if the notebook is executed in Colab. In
+    # case the notebook is executed in Jupyter notebook or other
+    # IPython runtimes, no interactive login is provided, it is
+    # assumed that the `GOOGLE_APPLICATION_CREDENTIALS` env var is
+    # set or `gcloud auth application-default login` was executed
+    # already. For more information on using Application Default Credentials
+    # see https://cloud.google.com/docs/authentication/production
+    # Attempt to connect to the Quantum Engine API, and use a simulator if unable to connect.
+    try:
+        from google.colab import auth
+    except ImportError:
+        print("Not running in a colab kernel. Will use Application Default Credentials.")
+        return
+
+    try:
+        print("Getting OAuth2 credentials.")
+        print("Press enter after entering the verification code.")
+        a = auth.authenticate_user(clear_output=clear_output)
+        print(a)
+        print("Authentication complete.")
+    except Exception as exc:
+        print(
+            "Authentication failed, you may not have permission to access a"
+            + " hardware Engine. Use a virtual Engine instead."
+        )
+        raise exc

cirq-google/cirq_google/engine/qcs_notebook_test.py

@@ -13,11 +13,16 @@
 # limitations under the License.
 
 import sys
+import unittest
 import unittest.mock as mock
 import pytest
 
 import cirq_google as cg
-from cirq_google.engine.qcs_notebook import get_qcs_objects_for_notebook, QCSObjectsForNotebook
+from cirq_google.engine.qcs_notebook import (
+    get_qcs_objects_for_notebook,
+    QCSObjectsForNotebook,
+    authenticate_user,
+)
 
 
 def _assert_correct_types(result: QCSObjectsForNotebook):
@@ -132,3 +137,30 @@ def test_get_qcs_objects_for_notebook_auth_fails(engine_mock):
     assert not result.signed_in
     assert result.is_simulator
     assert result.project_id == 'fake_project'
+
+
+class TestAuthenticateUser(unittest.TestCase):
+    """Tests for the public API `get_hardware_engine_and_authenticate_user` which
+    authenticates the user and returns a production engine instance ."""
+
+    @mock.patch.dict('sys.modules', {'google.colab': mock.Mock()})
+    def test_authentication_succeeds_no_exceptions_thrown(self):
+        authenticate_user("project_id")
+        auth_mock = sys.modules['google.colab']
+
+        assert auth_mock.auth.authenticate_user.called
+
+    @mock.patch.dict('sys.modules', {'google.colab': mock.Mock()})
+    def test_authentication_failure(self):
+        project_id = "invalid-project"
+        # Mock authentication failure
+        auth_mock = sys.modules['google.colab']
+
+        auth_mock.auth.authenticate_user = mock.Mock(side_effect=Exception('mock auth failure'))
+
+        with self.assertRaises(
+            Exception,
+            msg="Authentication failed, you may not have permission to access"
+            + " a hardware Engine. Use a virtual Engine instead.",
+        ):
+            authenticate_user(project_id)