def test_initialize_session_for_valid_regional_device(aws_session, caplog):
device_arn = f"arn:aws:braket:{aws_session.region}::device/qpu/test/device-name"
logger = logging.getLogger(__name__)
aws_session.get_device.return_value = device_arn
caplog.set_level(logging.INFO)
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
assert not caplog.text
def test_initialize_session_for_valid_non_regional_device(aws_session, caplog):
device_arn = "arn:aws:braket:::device/qpu/test/device-name"
first_region = aws_session.region
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = [
ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
),
ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
),
device_arn,
]
caplog.set_level(logging.INFO)
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
assert f"Changed session region from '{first_region}' to '{aws_session.region}'" in caplog.text
def test_initialize_session_local_device(mock_new_session, aws_session):
logger = logging.getLogger(__name__)
device = "local:provider.device.name"
# don't change a provided AwsSession
assert AwsQuantumJob._initialize_session(aws_session, device,
logger) == aws_session
# otherwise, create an AwsSession with the profile defaults
assert AwsQuantumJob._initialize_session(None, device,
logger) == mock_new_session()
def test_regional_device_raises_error(get_device_side_effect,
expected_exception, aws_session, caplog):
device_arn = f"arn:aws:braket:{aws_session.region}::device/qpu/test/device-name"
aws_session.get_device.side_effect = get_device_side_effect
logger = logging.getLogger(__name__)
caplog.set_level(logging.INFO)
with pytest.raises(expected_exception):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
aws_session.get_device.assert_called_with(device_arn)
assert not caplog.text
def test_bad_arn_format(aws_session):
logger = logging.getLogger(__name__)
device_not_found = (
"Device ARN is not a valid format: bad-arn-format. For valid Braket ARNs, "
"see 'https://docs.aws.amazon.com/braket/latest/developerguide/braket-devices.html'"
)
with pytest.raises(ValueError, match=device_not_found):
AwsQuantumJob._initialize_session(aws_session, "bad-arn-format",
logger)
def test_initialize_session_for_invalid_device(aws_session, device_arn):
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
)
device_not_found = "QPU 'arn:aws:braket:::device/qpu/test/device-name' not found."
with pytest.raises(ValueError, match=device_not_found):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_exception_in_credentials_session_region(device_arn, aws_session):
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = ClientError(
{"Error": {
"Code": "SomeOtherErrorMessage",
}},
"getDevice",
)
error_message = ("An error occurred \\(SomeOtherErrorMessage\\) "
"when calling the getDevice operation: Unknown")
with pytest.raises(ClientError, match=error_message):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_no_region_routing_simulator(aws_session):
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
)
device_arn = "arn:aws:braket:::device/simulator/test/device-name"
device_not_found = f"Simulator '{device_arn}' not found in 'us-test-1'"
with pytest.raises(ValueError, match=device_not_found):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_exceptions_in_all_device_regions(aws_session):
device_arn = "arn:aws:braket:::device/qpu/test/device-name"
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = [
ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
),
ClientError(
{"Error": {
"Code": "SomeOtherErrorMessage",
}},
"getDevice",
),
]
error_message = ("An error occurred \\(SomeOtherErrorMessage\\) "
"when calling the getDevice operation: Unknown")
with pytest.raises(ClientError, match=error_message):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_regional_device_switches(aws_session, caplog):
original_region = aws_session.region
device_region = "us-east-1"
device_arn = f"arn:aws:braket:{device_region}::device/qpu/test/device-name"
mock_session = Mock()
mock_session.get_device.side_effect = device_arn
aws_session.copy_session.side_effect = [mock_session]
logger = logging.getLogger(__name__)
caplog.set_level(logging.INFO)
assert mock_session == AwsQuantumJob._initialize_session(
aws_session, device_arn, logger)
aws_session.copy_session.assert_called_with(region=device_region)
mock_session.get_device.assert_called_with(device_arn)
assert f"Changed session region from '{original_region}' to '{device_region}'" in caplog.text