From ccee35604afc2b04d83ee9103eccb2821a4256cb Mon Sep 17 00:00:00 2001
From: Kshitij Chhabra <kshitijc@users.noreply.github.com>
Date: Fri, 1 Jul 2022 09:46:35 -0700
Subject: [PATCH] Add user agent for Braket interactions (#99)

---
 setup.py                                     |  2 +-
 src/braket/pennylane_plugin/braket_device.py |  2 ++
 test/unit_tests/test_braket_device.py        | 30 ++++++++++++++++----
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index 9393d0c1..32c9df1d 100644
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,7 @@
     packages=find_namespace_packages(where="src", exclude=("test",)),
     package_dir={"": "src"},
     install_requires=[
-        "amazon-braket-sdk",
+        "amazon-braket-sdk>=1.25.0",
         "pennylane>=0.23.0",
     ],
     entry_points={
diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py
index b5c70ca3..87b3402f 100644
--- a/src/braket/pennylane_plugin/braket_device.py
+++ b/src/braket/pennylane_plugin/braket_device.py
@@ -319,6 +319,8 @@ def __init__(
         **run_kwargs,
     ):
         device = AwsDevice(device_arn, aws_session=aws_session)
+        user_agent = f"BraketPennylanePlugin/{__version__}"
+        device.aws_session.add_braket_user_agent(user_agent)
         if DeviceActionType.JAQCD not in device.properties.action:
             raise ValueError(f"Device {device.name} does not support quantum circuits")
 
diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py
index 760b94d1..62d86ba8 100644
--- a/test/unit_tests/test_braket_device.py
+++ b/test/unit_tests/test_braket_device.py
@@ -34,7 +34,7 @@
 from pennylane.tape import QuantumTape
 
 import braket.pennylane_plugin.braket_device
-from braket.pennylane_plugin import BraketAwsQubitDevice, BraketLocalQubitDevice
+from braket.pennylane_plugin import BraketAwsQubitDevice, BraketLocalQubitDevice, __version__
 from braket.pennylane_plugin.braket_device import BraketQubitDevice, Shots
 
 SHOTS = 10000
@@ -753,6 +753,18 @@ def test_run_task_unimplemented():
     dev.execute(circuit)
 
 
+@patch("braket.pennylane_plugin.braket_device.AwsDevice")
+def test_add_braket_user_agent_invoked(aws_device_mock):
+    aws_device_mock_instance = aws_device_mock.return_value
+    aws_device_mock_instance.properties.action = {DeviceActionType.JAQCD: ACTION_PROPERTIES}
+    aws_device_mock_instance.type = AwsDeviceType.SIMULATOR
+    BraketAwsQubitDevice(wires=2, device_arn="foo", shots=10)
+    expected_user_agent = f"BraketPennylanePlugin/{__version__}"
+    aws_device_mock_instance.aws_session.add_braket_user_agent.assert_called_with(
+        expected_user_agent
+    )
+
+
 class DummyLocalQubitDevice(BraketQubitDevice):
     short_name = "dummy"
 
@@ -827,10 +839,17 @@ def _noop(*args, **kwargs):
 
 
 @patch.object(AwsDevice, "__init__", _noop)
+@patch.object(AwsDevice, "aws_session", new_callable=mock.PropertyMock)
 @patch.object(AwsDevice, "type", new_callable=mock.PropertyMock)
 @patch.object(AwsDevice, "properties")
 def _aws_device(
-    properties_mock, type_mock, wires, device_type=AwsDeviceType.QPU, shots=SHOTS, **kwargs
+    properties_mock,
+    type_mock,
+    session_mock,
+    wires,
+    device_type=AwsDeviceType.QPU,
+    shots=SHOTS,
+    **kwargs,
 ):
     properties_mock.action = {DeviceActionType.JAQCD: ACTION_PROPERTIES}
     properties_mock.return_value.action.return_value = {DeviceActionType.JAQCD: ACTION_PROPERTIES}
@@ -841,13 +860,14 @@ def _aws_device(
         device_arn="baz",
         aws_session=Mock(),
         shots=shots,
-        **kwargs
+        **kwargs,
     )
 
 
 @patch.object(AwsDevice, "__init__", _noop)
+@patch.object(AwsDevice, "aws_session", new_callable=mock.PropertyMock)
 @patch.object(AwsDevice, "properties")
-def _bad_aws_device(properties_mock, wires, **kwargs):
+def _bad_aws_device(properties_mock, session_mock, wires, **kwargs):
     properties_mock.action = {DeviceActionType.ANNEALING: ACTION_PROPERTIES}
     properties_mock.type = AwsDeviceType.QPU
     return BraketAwsQubitDevice(
@@ -856,5 +876,5 @@ def _bad_aws_device(properties_mock, wires, **kwargs):
         device_arn=DEVICE_ARN,
         aws_session=Mock(),
         shots=SHOTS,
-        **kwargs
+        **kwargs,
     )