CODEWITHSUNDEEP
pythonunit-testingmocking
Ankush
Ankush

Reputation: 1874

mock s3 connection and boto.S3key to check set_content_from_string method

I am doing unit test with python mock. I've gone through blogs and python docs related to mocking but confuse about mocking the test case. Here is the snippet for which I want to write test case. The agenda is to test the method "set_contents_from_string()" using mock.

def write_to_customer_registry(customer):

    # establish a connection with S3
    conn = _connect_to_s3()

    # build customer registry dict and convert it to json
    customer_registry_dict = json.dumps(build_customer_registry_dict(customer))

    # attempt to access requested bucket
    bucket = _get_customer_bucket(conn)

    s3_key = _get_customer_key(bucket, customer)
    s3_key.set_metadata('Content-Type', 'application/json')
    s3_key.set_contents_from_string(customer_registry_dict)
    return s3_key

Upvotes: 0

Views: 2202

Answers (2)

Ankush
Ankush

Reputation: 1874

came up with solution that worked for me, Posting it here, may be helpful for someone.

def setup(self):
    self.customer = Customer.objects.create('tiertranstests')
    self.customer.save()

def test_build_customer_registry(self):
    mock_connection = Mock()
    mock_bucket = Mock()
    mock_s3_key = Mock()

    customer_registry_dict = json.dumps(build_customer_registry_dict(self.customer))
    # Patch S3 connection and Key class of registry method
    with patch('<path>.customer_registry.S3Connection', Mock(return_value=mock_connection)),\
         patch('<path>.customer_registry.Key', Mock(return_value=mock_s3_key)):

        mock_connection.get_bucket = Mock(return_value=mock_bucket)
        mock_s3_key.set_metadata.return_value = None
        mock_s3_key.set_contents_from_string = Mock(return_value=customer_registry_dict)

        write_to_customer_registry(self.customer)
        mock_s3_key.set_contents_from_string.assert_called_once_with(customer_registry_dict)

Upvotes: 1

Enrique Saez
Enrique Saez

Reputation: 2700

As you are testing some private methods I have added them to a module which I called s3.py that contains your code:

import json

def _connect_to_s3():
    raise

def _get_customer_bucket(conn):
    raise

def _get_customer_key(bucket, customer):
    raise

def build_customer_registry_dict(cust):
    raise

def write_to_customer_registry(customer):

    # establish a connection with S3
    conn = _connect_to_s3()

    # build customer registry dict and convert it to json
    customer_registry_dict = json.dumps(build_customer_registry_dict(customer))

    # attempt to access requested bucket
    bucket = _get_customer_bucket(conn)

    s3_key = _get_customer_key(bucket, customer)
    s3_key.set_metadata('Content-Type', 'application/json')
    s3_key.set_contents_from_string(customer_registry_dict)
    return s3_key

Next, in another module test_s3.py, I tested your code taking into account that for Unit Tests all interactions with third parties, such as network calls to s3 should be patched:

from unittest.mock import MagicMock, Mock, patch
from s3 import write_to_customer_registry
import json

@patch('json.dumps', return_value={})
@patch('s3._get_customer_key')
@patch('s3.build_customer_registry_dict')
@patch('s3._get_customer_bucket')
@patch('s3._connect_to_s3')
def test_write_to_customer_registry(connect_mock, get_bucket_mock, build_customer_registry_dict_mock, get_customer_key_mock, json_mock):

  customer = MagicMock()
  connect_mock.return_value = 'connection'
  get_bucket_mock.return_value = 'bucket'
  get_customer_key_mock.return_value = MagicMock()

  write_to_customer_registry(customer)

  assert connect_mock.call_count == 1
  assert get_bucket_mock.call_count == 1
  assert get_customer_key_mock.call_count == 1

  get_bucket_mock.assert_called_with('connection')
  get_customer_key_mock.assert_called_with('bucket', customer)
  get_customer_key_mock.return_value.set_metadata.assert_called_with('Content-Type', 'application/json')
  get_customer_key_mock.return_value.set_contents_from_string.assert_called_with({})

As you can see from the tests I am not testing that set_contents_from_string is doing what is supposed to do (since that should already be tested by the boto library) but that is being called with the proper arguments. If you still doubt that the boto library is not properly testing such call you can always check it yourself in boto Github or boto3 Github

Something else you could test is that your are handling the different exceptions and edge cases in your code properly.

Finally, you can find more about patching and mocking in the docs. Usually the section about where to patch is really useful.

Some other resources are this blog post with python mock gotchas or this blog post I wrote myself (shameless self plug) after answering related pytest, patching and mocking questions in Stackoverflow.

Upvotes: 1

Related Questions