test.py

#!/usr/bin/env python3.7

import boto3
import botoinator
from moto import mock_s3, mock_sqs
import time
import inspect

""" This is our decorator that we will apply to boto3 methods """
def myDecorator():
  def test_decorator(func):
    print('fired')
    # Add an attribute for testing
    func.testValue = True
    return func
  return test_decorator

@mock_s3
def testRegisterToClient():

  """
  Test registering a decorator to a single boto3 session
  """

  # Create a boto3 session
  s = boto3.session.Session()

  # Register the create_bucket() method to use our decorator for this session
  s.register_client_decorator('s3', 'create_bucket', myDecorator())

  # Now create our client as we normally would
  client = s.client('s3')

  # Now we can see that create_bucket() was decorated by testing the attribute we added
  assert hasattr(client.create_bucket, 'testValue')

  # We can also see that this only applies to calls made by the session we registered by creating a new session through boto3.client() and not registering a decorator
  client = boto3.client('s3')

  # Now we can see that client.create_bucket() is not decorated
  assert not hasattr(client.create_bucket, 'testValue')


@mock_sqs
def testRegisterToResource():

  """
  Test registering a decorator to a single boto3 session
  """

  # Create a boto3 session
  s = boto3.session.Session()

  # Register the delete() method of SQS.Queue resource to be decorated for this session
  s.register_resource_decorator('sqs', 'Queue', 'delete', myDecorator())

  # Create our resource
  sqs = s.resource('sqs', region_name='us-east-1')

  # Create a queue to test with
  q = sqs.create_queue(QueueName='foo')

  # Now we can see that our decorator was called by testing the testValue attribute to SQS.Queue.delete() method
  assert hasattr(sqs.Queue('foo').delete, 'testValue')

  # We can also see that this only applies to calls made by the session we registered by creating a new session through boto3.resource() and not registering a decorator
  sqs = boto3.resource('sqs', region_name='us-east-1')

  # Now we can see that SQS.Queue.delete() is not decorated
  assert not hasattr(sqs.Queue('foo').delete, 'testValue')

@mock_s3
def testAddToClient():

  """
  Test registering a decorator to all boto3 sessions created
  """

  # Register the create_bucket() method to use our decorator with some args for testing
  boto3.session.Session.add_client_decorator('s3', 'create_bucket', myDecorator())

  # Now we can see that create_bucket() was decorated for two different clients/sessions by testing the attribute we added
  clientA = boto3.client('s3')
  assert hasattr(clientA.create_bucket, 'testValue')

  clientB = boto3.client('s3')
  assert hasattr(clientB.create_bucket, 'testValue')


@mock_sqs
def testAddToResource():

  """
  Test registering a decorator to all boto3 sessions created
  """

  # Register the create_bucket() method to use our decorator with some args for testing
  boto3.session.Session.add_resource_decorator('sqs', 'Queue', 'delete', myDecorator())


  # Now create two clients
  sqs1 = boto3.resource('sqs', region_name='us-east-1')
  sqs2 = boto3.resource('sqs', region_name='us-east-1')

  # Create a queue to test with
  sqs1.create_queue(QueueName='foo')

  # Now we can see that our decorator was called by testing the testValue attribute to SQS.Queue.delete() method
  assert hasattr(sqs1.Queue('foo').delete, 'testValue')
  assert hasattr(sqs2.Queue('foo').delete, 'testValue')

#testRegisterToClient()
#testRegisterToResource()
testAddToClient()
testAddToResource()