"""Test aws utils sqs.""" import boto3 import pytest from moto import mock_aws from vector_utils.aws_utils import sqs from vector_utils.exceptions import ValidationError @pytest.mark.parametrize( 'env,encoder_id,priority,dms_priority,expected', [ ('dev', 1, 1, 1, 'dev-delivery1_e0000001_d0000001'), ('dev', 100, 100, 100, 'dev-delivery100_e0000100_d0000100') ]) def test_delivery_queue_url(env, encoder_id, priority, dms_priority, expected): """Test that we generate queue url properly.""" result = sqs.format_queue_name( sqs.DELIVERY_QUEUE_NAME_PATTERN, env=env, encoder_id=encoder_id, priority=priority, dms_priority=dms_priority) assert result == expected @pytest.mark.parametrize( 'env,expected', [ ('dev', 'dev-batch_to_close_v2-queue'), ('prod', 'prod-batch_to_close_v2-queue') ]) def test_close_batch_queue_url(env, expected): """Test that we generate close batch queue url properly.""" result = sqs.format_queue_name( sqs.BATCH_TO_CLOSE_QUEUE_NAME_PATTERN, env=env) assert result == expected @mock_aws def test_add_delivery_job_to_sqs(): """Test add_delivery_job_to_sqs.""" queue_name = 'dev-delivery1_e0000001_d0000001' job_data = sqs.add_message_to_sqs(queue_name, { 'eqd_id': 1, 'dms_id': 1, 'delivery_batch_id': 1 }) assert job_data.keys() & {'Successful', 'Failed'} @mock_aws def test_get_messages_from_sqs(): """Test get_messages_from_sqs.""" queue_name = 'dev-batch_to_close_v2-queue' sqs_client = boto3.client('sqs', region_name='us-east-1') sqs_client.create_queue(QueueName=queue_name) sqs.add_message_to_sqs(queue_name, { 'eqd_id': 1, 'dms_id': 1, 'delivery_batch_id': 1 }) jobs = sqs.get_messages_from_sqs(queue_name) expected = '{"eqd_id": 1, "dms_id": 1, "delivery_batch_id": 1}' assert jobs[0]['Body'] == expected @mock_aws def test_get_messages_from_sqs_non_exist_queue(): """Test get_messages_from_sqs.""" queue_name = 'dev-batch_to_close_v2-queue' jobs = sqs.get_messages_from_sqs(queue_name) expected = [] assert jobs == expected @mock_aws def test_delete_messages_from_sqs(): """Test get_messages_from_sqs.""" queue_name = 'dev-batch_to_close_v2-queue' sqs_client = boto3.client('sqs', region_name='us-east-1') sqs_client.create_queue(QueueName=queue_name) sqs.add_message_to_sqs(queue_name, { 'eqd_id': 1, 'dms_id': 1, 'delivery_batch_id': 1 }) jobs = sqs.get_messages_from_sqs(queue_name) receipts = [job['ReceiptHandle'] for job in jobs] deleted = sqs.delete_messages_from_sqs(queue_name, receipts) assert len(deleted) == 1 @mock_aws def test_add_messages_to_sqs(): """Test add_messages_to_sqs.""" queue_name = 'dev-delivery1_e0000001_d0000001' results = sqs.add_messages_to_sqs( queue_name, [ { 'eqd_id': 1, 'dms_id': 1, 'delivery_batch_id': 1 }, 'message body', ] ) assert results.keys() & {'Successful', 'Failed'} assert len(results.get('Successful')) == 2 assert len(results.get('Failed')) == 0 @pytest.mark.parametrize( 'queue_name_test_pattern,env,encoder_id,dms_priority,priority', [ ('{env}-delivery{encoder_id}_e{priority:07d}_d{dms_priority:07d}', 'dev', 1, 1, 1), ('{env}-encoding{encoder_id}_e{priority:07d}_d{dms_priority:07d}', 'dev', 1, 1, 1), ]) def test_format_queue_name_delivery_and_encoding( queue_name_test_pattern, env, encoder_id, dms_priority, priority): """Test for formatting the queue names. delivery encoding """ expected = queue_name_test_pattern.format( env=env, encoder_id=encoder_id, priority=priority, dms_priority=dms_priority) result = sqs.format_queue_name( queue_name_test_pattern, env=env, encoder_id=encoder_id, dms_priority=dms_priority, priority=priority) assert result == expected @pytest.mark.parametrize( 'queue_name_test_pattern,env,encoder_id,dms_priority,priority,loco_id', [ ('{env}-batch_to_close_v2-queue', 'dev', None, None, None, None), ('{env}-assets_to_download-queue', 'dev', None, None, None, None), ('{env}-assets_to_copy_d{loco_id:07d}-queue', 'dev', None, None, None, 1), ]) def test_format_queue_name_batch_download_copy( queue_name_test_pattern, env, encoder_id, dms_priority, priority, loco_id): """Test for formatting queue names. batch_to_close assets_to_download assets_to_copy """ expected = queue_name_test_pattern.format( env=env, encoder_id=encoder_id, priority=priority, dms_priority=dms_priority, loco_id=loco_id) result = sqs.format_queue_name( queue_name_test_pattern, env=env, encoder_id=encoder_id, dms_priority=dms_priority, priority=priority, loco_id=loco_id) assert result == expected def test_format_queue_name_assets_to_copy_raises_validation_error(): """Test that format_queue_name for assets_to_copy pattern. raises ValidationError """ queue_name_test_pattern = '{env}-assets_to_copy_d{loco_id:07d}-queue' with pytest.raises(ValidationError): sqs.format_queue_name( queue_name_test_pattern, env=None, encoder_id=None, dms_priority=None, priority=None, loco_id=None) @pytest.mark.parametrize( 'queue_name_test_pattern,env,encoder_id,dms_priority,priority', [ ('{env}-delivery{encoder_id}_e{priority:07d}_d{dms_priority:07d}', 'dev', None, None, None), ('{env}-delivery{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, 1, None, None), ('{env}-delivery{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, None, 1, None), ('{env}-delivery{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, None, None, 1), ('{env}-delivery{encoder_id}_e{priority:07d}_d{dms_priority:07d}', 'dev', None, None, None), ('{env}-encoding{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, 1, None, None), ('{env}-encoding{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, None, 1, None), ('{env}-encoding{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, None, None, 1), ('{env}-encoding{encoder_id}_e{priority:07d}_d{dms_priority:07d}', None, None, None, None), ]) def test_format_queue_name_delivery_and_encoding_raises_validation_error( queue_name_test_pattern, env, encoder_id, dms_priority, priority): """Test that format_queue_name for assets_to_copy pattern. raises ValidationError """ with pytest.raises(ValidationError): sqs.format_queue_name( queue_name_test_pattern, env=env, encoder_id=encoder_id, dms_priority=dms_priority, priority=priority)