Using Mock in Python for nested objects (DynamoDB and Table)

Question

I want to test the behavior on a function with different inputs from a DynamoDb. There are two main behaviors: when the search key is found in the table and when it is not. Here is a minimal code of the function:

import boto3
from boto3.dynamodb.conditions import Key

def main(symbol):
   dynamo = boto3.resource("dynamodb")
   table = dynamo.Table("mytable")
   data = table.query(KeyConditionExpression=Key("symbol").eq(symbol))

   if data.count > 0:
      # result = some_output
   else:
      result = {'status': '404'}
   return result

and I want to test this code with a unit test by sending empty results and a list of non-empty items, something along these lines:

import boto3
import unittest
from unittest.mock import Mock, patch

class TestMainHandler(unittest.TestCase): 
   ...
   def test_main_fails_on_wrong_symbol(self):
       with patch.object(main_handler, 'table') as get_mock:
          get_mock.return_value = []
          result = main('dummy_symbol')
          expect_result = {'status': '404'}
   self.assertEqual(result, expect_result)

but I can't run the mock part. I was wondering if you could guide me how to mock the nested table and dynamo variables. I highly appreciate your kind help.

Answer 1

I'd recommend mocking the first "integration point" you can, then building your mock results off of that. In this case, it would be boto3.resource. From there, you can modify the return value of boto3.resource to be a mock table. Then you can change the return value of any call on your mock table to be your expected result.

import boto3
import unittest
from unittest.mock import Mock, patch

class TestMainHandler(unittest.TestCase): 
    @patch('boto3.resource')
    def test_main_fails_on_wrong_symbol(self, mock_dynamo):
        mock_table = Mock()
        mock_table.query.return_value = []
        mock_dynamo.Table.return_value = mock_table

        result = main('dummy_symbol')
        expected_result = {'status': '404'}
        self.assertEqual(expected_result, result)

Also note that I used the patch decorator on the test case rather than the context manager. This is a matter of preference for me, but I think it looks cleaner.

EDIT: I had an error when assigning the return value for the mock_dynamo function call. I also filled in your if to test the length of the queried data and return something. That branch is not reached by the test here though. Here's the final product:

"""boto_main.py"""
import boto3
from boto3.dynamodb.conditions import Key

def main(symbol):
   dynamo = boto3.resource("dynamodb")
   table = dynamo.Table("mytable")
   data = table.query(KeyConditionExpression=Key("symbol").eq(symbol))
   if len(data) > 0:
      result = {'status': '200'}
   else:
      result = {'status': '404'}
   return result

---

"""boto_test.py"""
import unittest
from unittest.mock import Mock, patch
from boto_main import main

class TestMainHandler(unittest.TestCase): 
    @patch('boto3.resource')
    def test_main_fails_on_wrong_symbol(self, mock_dynamo):
        mock_table = Mock()
        mock_table.query.return_value = []
        mock_dynamo.return_value.Table.return_value = mock_table
        result = main('dummy_symbol')
        expected_result = {'status': '404'}
        self.assertEqual(expected_result, result)