Aws device fix (#161)

fix: better exception handling for aws device
amazon-braket · Nov 19, 2021 · a90cf35 · a90cf35
1 parent 7a2c933
commit a90cf35
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 2 deletions.
diff --git a/src/braket/aws/aws_device.py b/src/braket/aws/aws_device.py
@@ -244,8 +244,9 @@ def _get_session_and_initialize(self, session):
             try:
                 self._populate_properties(region_session)
                 return region_session
-            except Exception:
-                pass
+            except ClientError as e:
+                if e.response["Error"]["Code"] != "ResourceNotFoundException":
+                    raise e
         raise ValueError(f"QPU '{self._arn}' not found")
 
     def _populate_properties(self, session):

diff --git a/test/unit_tests/braket/aws/test_aws_device.py b/test/unit_tests/braket/aws/test_aws_device.py
@@ -433,6 +433,44 @@ def test_device_qpu_not_found(mock_copy_session):
         AwsDevice("arn:aws:braket:::device/qpu/a/b", mock_session)
 
 
+@patch("braket.aws.aws_device.AwsSession.copy_session")
+def test_device_qpu_exception(mock_copy_session):
+    mock_session = Mock()
+    mock_session.get_device.side_effect = (
+        ClientError(
+            {
+                "Error": {
+                    "Code": "ResourceNotFoundException",
+                    "Message": (
+                        "Braket device 'arn:aws:braket:::device/quantum-simulator/amazon/tn1' "
+                        "not found in us-west-1. You can find a list of all supported device "
+                        "ARNs and the regions in which they are available in the documentation: "
+                        "https://docs.aws.amazon.com/braket/latest/developerguide/braket-"
+                        "devices.html"
+                    ),
+                }
+            },
+            "getDevice",
+        ),
+        ClientError(
+            {
+                "Error": {
+                    "Code": "OtherException",
+                    "Message": "Some other message",
+                }
+            },
+            "getDevice",
+        ),
+    )
+    mock_copy_session.return_value = mock_session
+    qpu_exception = (
+        "An error occurred \\(OtherException\\) when calling the "
+        "getDevice operation: Some other message"
+    )
+    with pytest.raises(ClientError, match=qpu_exception):
+        AwsDevice("arn:aws:braket:::device/qpu/a/b", mock_session)
+
+
 @patch("braket.aws.aws_device.AwsSession.copy_session")
 def test_device_non_qpu_region_error(mock_copy_session):
     mock_session = Mock()