diff --git a/src/braket/aws/queue_information.py b/src/braket/aws/queue_information.py index d0a105a31..458b3f3f8 100644 --- a/src/braket/aws/queue_information.py +++ b/src/braket/aws/queue_information.py @@ -35,7 +35,7 @@ class QueueDepthInfo: Represents quantum tasks and hybrid jobs queue depth information. Attributes: - quantum_task (str, Dict[QueuePriority, str]): number of quantum_tasks waiting + quantum_task (Dict[QueueType, str]): number of quantum_tasks waiting to run on a device. This includes both 'Normal' and 'Priority' tasks. For Example, {'quantum_task': {QueueType.NORMAL: '7', QueueType.PRIORITY: '3'}} job (str): number of hybrid jobs waiting to run on a device. Additionally, for QPUs if diff --git a/test/integ_tests/test_queue_position.py b/test/integ_tests/test_queue_information.py similarity index 77% rename from test/integ_tests/test_queue_position.py rename to test/integ_tests/test_queue_information.py index 1292e85f5..6c1401990 100644 --- a/test/integ_tests/test_queue_position.py +++ b/test/integ_tests/test_queue_information.py @@ -11,9 +11,13 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. - from braket.aws import AwsDevice, AwsQuantumJob -from braket.aws.queue_information import HybridJobQueueInfo, QuantumTaskQueueInfo, QueueType +from braket.aws.queue_information import ( + HybridJobQueueInfo, + QuantumTaskQueueInfo, + QueueDepthInfo, + QueueType, +) from braket.circuits import Circuit from braket.devices import Devices @@ -62,3 +66,19 @@ def test_job_queue_position(aws_session): # assert message assert queue_information.queue_position is None assert isinstance(queue_information.message, str) + + +def test_queue_depth(): + device = AwsDevice(Devices.Amazon.SV1) + + # call the queue_depth method. + queue_information = device.queue_depth() + + # data type validations + assert isinstance(queue_information, QueueDepthInfo) + assert isinstance(queue_information.quantum_task, dict) + assert isinstance(queue_information.job, str) + + for key, value in queue_information.quantum_task.items(): + assert isinstance(key, QueueType) + assert isinstance(value, str)