lambda.txt

import pytest
from moto import mock_dynamodb2, mock_sqs
import boto3
import os
import json
from lambdas.scanner.scanner_lambda import lambda_handler

@pytest.fixture
def aws_credentials():
    os.environ['AWS_ACCESS_KEY_ID'] = 'testing'
    os.environ['AWS_SECRET_ACCESS_KEY'] = 'testing'
    os.environ['AWS_SECURITY_TOKEN'] = 'testing'
    os.environ['AWS_SESSION_TOKEN'] = 'testing'

@pytest.fixture
def dynamodb(aws_credentials):
    with mock_dynamodb2():
        yield boto3.resource('dynamodb')

@pytest.fixture
def sqs(aws_credentials):
    with mock_sqs():
        yield boto3.client('sqs')

def test_scanner_lambda_single_page(dynamodb, sqs):
    # Create DynamoDB table
    table = dynamodb.create_table(
        TableName='ContractTable',
        KeySchema=[
            {'AttributeName': 'contract_id', 'KeyType': 'HASH'},
            {'AttributeName': 'instrument_class_id', 'KeyType': 'RANGE'}
        ],
        AttributeDefinitions=[
            {'AttributeName': 'contract_id', 'AttributeType': 'S'},
            {'AttributeName': 'instrument_class_id', 'AttributeType': 'S'}
        ],
        ProvisionedThroughput={'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1}
    )
    # Add items
    table.put_item(Item={'contract_id': '1', 'instrument_class_id': 'A'})
    table.put_item(Item={'contract_id': '2', 'instrument_class_id': 'B'})

    # Create SQS queue
    queue_url = sqs.create_queue(QueueName='DeletionQueue')['QueueUrl']

    # Set environment variables
    os.environ['TABLE_NAME'] = 'ContractTable'
    os.environ['QUEUE_URL'] = queue_url

    # Define event
    event = {
        'segment': 0,
        'total_segments': 1,
        'instrument_class_id': 'A'
    }

    # Call Lambda handler
    response = lambda_handler(event, None)

    # Verify SQS messages
    messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)
    assert len(messages['Messages']) == 1
    assert json.loads(messages['Messages'][0]['Body']) == '1'

    # Verify no pagination
    assert 'LastEvaluatedKey' not in response