From f1a549a678a057c5c1341d5099a7e7cdac7db333 Mon Sep 17 00:00:00 2001
From: Milan <30416311+krneta@users.noreply.github.com>
Date: Wed, 15 Nov 2023 10:00:39 -0800
Subject: [PATCH] update: default no longer returning RETIRED devices from
 get_devices (#796)

---
 src/braket/aws/aws_device.py                  | 12 +++++++-----
 src/braket/aws/aws_session.py                 | 12 ++++++++----
 test/integ_tests/test_device_creation.py      |  2 +-
 .../unit_tests/braket/aws/test_aws_session.py | 19 +++++++++++++++++++
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/braket/aws/aws_device.py b/src/braket/aws/aws_device.py
index c9a208411..0cff446db 100644
--- a/src/braket/aws/aws_device.py
+++ b/src/braket/aws/aws_device.py
@@ -564,13 +564,15 @@ def get_devices(
             >>> AwsDevice.get_devices(types=['SIMULATOR'])
 
         Args:
-            arns (Optional[list[str]]): device ARN list, default is `None`
-            names (Optional[list[str]]): device name list, default is `None`
-            types (Optional[list[AwsDeviceType]]): device type list, default is `None`
+            arns (Optional[list[str]]): device ARN filter, default is `None`
+            names (Optional[list[str]]): device name filter, default is `None`
+            types (Optional[list[AwsDeviceType]]): device type filter, default is `None`
                 QPUs will be searched for all regions and simulators will only be
                 searched for the region of the current session.
-            statuses (Optional[list[str]]): device status list, default is `None`
-            provider_names (Optional[list[str]]): provider name list, default is `None`
+            statuses (Optional[list[str]]): device status filter, default is `None`. When `None`
+                is used, RETIRED devices will not be returned. To include RETIRED devices in
+                the results, use a filter that includes "RETIRED" for this parameter.
+            provider_names (Optional[list[str]]): provider name filter, default is `None`
             order_by (str): field to order result by, default is `name`.
                 Accepted values are ['arn', 'name', 'type', 'provider_name', 'status']
             aws_session (Optional[AwsSession]): An AWS session object.
diff --git a/src/braket/aws/aws_session.py b/src/braket/aws/aws_session.py
index 9523f0573..2ebc3c8d2 100644
--- a/src/braket/aws/aws_session.py
+++ b/src/braket/aws/aws_session.py
@@ -622,10 +622,12 @@ def search_devices(
         all the filters `arns`, `names`, `types`, `statuses`, `provider_names`.
 
         Args:
-            arns (Optional[list[str]]): device ARN list, default is `None`.
-            names (Optional[list[str]]): device name list, default is `None`.
-            types (Optional[list[str]]): device type list, default is `None`.
-            statuses (Optional[list[str]]): device status list, default is `None`.
+            arns (Optional[list[str]]): device ARN filter, default is `None`.
+            names (Optional[list[str]]): device name filter, default is `None`.
+            types (Optional[list[str]]): device type filter, default is `None`.
+            statuses (Optional[list[str]]): device status filter, default is `None`. When `None`
+                is used, RETIRED devices will not be returned. To include RETIRED devices in
+                the results, use a filter that includes "RETIRED" for this parameter.
             provider_names (Optional[list[str]]): provider name list, default is `None`.
 
         Returns:
@@ -645,6 +647,8 @@ def search_devices(
                     continue
                 if statuses and result["deviceStatus"] not in statuses:
                     continue
+                if statuses is None and result["deviceStatus"] == "RETIRED":
+                    continue
                 if provider_names and result["providerName"] not in provider_names:
                     continue
                 results.append(result)
diff --git a/test/integ_tests/test_device_creation.py b/test/integ_tests/test_device_creation.py
index decd7d876..4cb7de2b1 100644
--- a/test/integ_tests/test_device_creation.py
+++ b/test/integ_tests/test_device_creation.py
@@ -18,7 +18,7 @@
 from braket.aws import AwsDevice
 from braket.devices import Devices
 
-RIGETTI_ARN = "arn:aws:braket:::device/qpu/rigetti/Aspen-10"
+RIGETTI_ARN = "arn:aws:braket:us-west-1::device/qpu/rigetti/Aspen-M-3"
 IONQ_ARN = "arn:aws:braket:us-east-1::device/qpu/ionq/Harmony"
 SIMULATOR_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
 OQC_ARN = "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy"
diff --git a/test/unit_tests/braket/aws/test_aws_session.py b/test/unit_tests/braket/aws/test_aws_session.py
index d2ee45a3a..c61d22606 100644
--- a/test/unit_tests/braket/aws/test_aws_session.py
+++ b/test/unit_tests/braket/aws/test_aws_session.py
@@ -673,6 +673,18 @@ def test_cancel_job_surfaces_errors(exception_type, aws_session):
                 },
             ],
         ),
+        (
+            {"statuses": ["RETIRED"]},
+            [
+                {
+                    "deviceArn": "arn4",
+                    "deviceName": "name4",
+                    "deviceType": "QPU",
+                    "deviceStatus": "RETIRED",
+                    "providerName": "pname3",
+                },
+            ],
+        ),
         (
             {"provider_names": ["pname2"]},
             [
@@ -745,6 +757,13 @@ def test_search_devices(input, output, aws_session):
                     "deviceStatus": "ONLINE",
                     "providerName": "pname2",
                 },
+                {
+                    "deviceArn": "arn4",
+                    "deviceName": "name4",
+                    "deviceType": "QPU",
+                    "deviceStatus": "RETIRED",
+                    "providerName": "pname3",
+                },
             ]
         }
     ]