Python Mocking Examples

Here the first example using decorator in python 2.7

import unittest
import random
import mock

def myrandom(p):
    return random.random() > p

class Test(unittest.TestCase):

    @mock.patch('random.random')
    def test_myrandom(self, mock_random):
        mock_random.return_value = 0.1
        val = myrandom(0.0)
        assert val  > 0
        assert mock_random.call_count == 1

if __name__ == '__main__':
    unittest.main()

Here the example for assert_callled_with() function

import unittest
import mock
import example

class Test(unittest.TestCase):

    @mock.patch('example.hello')
    def test1(self,mock_hello):
        x = 'Oj'
        example.hello(x)  # Uses patched example.func
        mock_hello.assert_called_with(x)

if __name__ == '__main__':
    unittest.main()

Above test can be ran using context manager:

import unittest
import mock
import example

class Test(unittest.TestCase):

    def test1(self):
        x = "OJ"
        with mock.patch('example.hello') as mock_hello:
            example.hello(x)
            mock_hello.assert_called_with(x)

if __name__ == '__main__':
    unittest.main()

If you are testing the method in the same module:

import unittest
import mock

def hello(name):
    return '{}'.formt(name)

class Test(unittest.TestCase):

    def test1(self):
        x = "OJ"
        with mock.patch('__main__.hello') as mock_hello:
            hello(x)
            mock_hello.assert_called_with(x)

if __name__ == '__main__':
    unittest.main()

This is how return value is working

import unittest
import mock

def hello(name):
    return None

class Test(unittest.TestCase):

    @mock.patch('__main__.hello', return_value='Oj')
    def test1(self,mock_hello):
        x = 'Oj'
        assert mock_hello(x) == x


if __name__ == '__main__':
    unittest.main()

If you run in the context manager:

import unittest
import mock

def hello(name):
    return None

class Test(unittest.TestCase):

    def test1(self):
        x = "OJ"
        with mock.patch('__main__.hello', return_value='Hello {}'.format(x)) as mock_hello:
            assert hello(x) == 'Hello {}'.format(x)

if __name__ == '__main__':
    unittest.main()

How to mock a method

import unittest
import mock

class MyMock(object):
    def __init__(self,name):
        self.name = name

attrs = {'method.return_value':MyMock('Tixtrac'), 'other.side_effect':KeyError}
mock = mock.Mock(some_attribute = 'eggs', **attrs)
mock.some_attribute
obj = mock.method()
print(obj.name)
mock.other()

In the above source, the method has been mocked.

Mock the class method, here because spec=True , you can only mock the methods belongs to SomeClass.

import unittest
import mock

class SomeClass(object):
    def hello():
        return 'Hello'

class Test(unittest.TestCase):
    @mock.patch('__main__.SomeClass', spec=True)
    def test_SomeClass(self, mc):

        #mock the method
        instance = mc.return_value
        instance.hello.return_value = 'Foo'
        #instance.m.return_value = 'Foo' #failed because spec=True

        #create instance of SomeClass
        sc = SomeClass()

        assert sc is instance
        assert sc.hello() == 'Foo'
        # assert sc.m() == 'Foo' #faiedl because spec=True

if __name__ == '__main__':
    unittest.main()

Here the way to mock the property decorator:

import unittest
import mock

class Foo(object):
    @property
    def foo(self):
        return 'foo thing'
    @foo.setter
    def foo(self, value):
        pass

class Test(unittest.TestCase):

    @mock.patch('__main__.Foo.foo', new_callable=mock.PropertyMock)
    def test1(self,mock):
        mock.return_value = 'test'
        foo = Foo()
        print('val: {}'.format(foo.foo))


if __name__ == '__main__':
    unittest.main()

Use alternative class for the default MagicMock using new_callable:

import unittest
import mock
from StringIO import StringIO
def foo():
    print 'Hello'

class Test(unittest.TestCase):
    @mock.patch('sys.stdout', new_callable=StringIO)
    def test_foo(self, mock): #mocked the stdout here
        foo() #first call the method
        assert mock.getvalue() == 'Hello\n'

if __name__ == '__main__':
    unittest.main()

Patch object

This can be used as:

  • decorator
  • class decorator
  • context manager

Patch object is either two or three arguments.

  1. object to be patched
  2. attribute name
  3. object to replace the attribute with

Here the example:

import unittest
import mock


class SomeClass(object):
    def hello(name):
        return '{} world!'.format(name)


class Test(unittest.TestCase):
    @mock.patch.object(SomeClass, 'hello')
    def test_SomeClass(self, mc):
        hello_str = 'Oj'
        SomeClass.hello(hello_str)
        mc.assert_called_with(hello_str)

if __name__ == '__main__':
    unittest.main()

Mocking the dictionary:

import unittest
import mock
import os

class Test(unittest.TestCase):

    @mock.patch.dict(os.environ, {'TARGET_STREAM' : 'test'})
    def test_dict(self):
        assert os.environ['TARGET_STREAM'] == 'test'

if __name__ == '__main__':
    unittest.main()

For the existing dictionary:

import unittest
import mock
import os
foo = {}
class Test(unittest.TestCase):

    @mock.patch.dict(foo, {'TARGET_STREAM' : 'test'})
    def test_dict(self):
        assert foo['TARGET_STREAM'] == 'test'

if __name__ == '__main__':
    unittest.main()

Mock the module using dictionary

import unittest
import mock
import os

mymodule = mock.MagicMock()
mymodule.hello.return_value = 'Hello'
class Test(unittest.TestCase):

    @mock.patch.dict('sys.modules', mymodule=mymodule)
    def test_dict(self):
        import mymodule #import this to test
        assert mymodule.hello('2') == 'Hello'

if __name__ == '__main__':
    unittest.main()

Tips to Mock AWS Boto 3

First create the sharerable context:

from contextlib import contextmanager
import os

@contextmanager
def lambdaEnv(stream_name):
    STREAM_NAME = stream_name
    os.environ['LOGGING_LEVEL'] = 'DEBUG'
    os.environ['SOURCE_STREAM'] = STREAM_NAME
    os.environ['TARGET_BUCKET'] = 'ojithak'
    os.environ['TARGET_BUCKET_ROOT_DIR'] = 'ticketingkin'

    os.environ['KS_CUSTOMER_SALES'] = 'arn:aws:kinesis:ap-southeast-2:167856709666:stream/customersales_in'#only use this
    os.environ['RT_SALES'] = 'sales'
    os.environ['KS_STORM'] = 'arn:aws:kinesis:ap-southeast-2:167856709666:stream/storm-dev'
    os.environ['RT_PAYMENT_CT'] = 'Sales.cdc.dbo_Payment_CT'
    os.environ['RT_ORDERLINES'] = 'OrdersQA.Orderlines'
    os.environ['RT_ORDERLINE_PAYMENTS'] = 'OrdersQA.Orderlines'
    os.environ['KS_TICKETING'] = 'arn:aws:kinesis:ap-southeast-2:167856709666:stream/ticketing-dev'
    os.environ['KS_ORDERS'] = 'arn:aws:kinesis:ap-southeast-2:167856709666:stream/new-orders-dev'
    os.environ['KS_TIXTRACK'] = 'arn:aws:kinesis:ap-southeast-2:167856709666:stream/tixtrack_dev'

    logging.basicConfig(format='[%(asctime)s] p%(process)s {%(filename)s:%(funcName)s:%(lineno)d} %(levelname)s - %(message)s',
                        datefmt='%d-%m-%Y:%H:%M:%S',
                        level=os.environ['LOGGING_LEVEL'])
    ch = logging.StreamHandler()
    log = logging.getLogger()
    log.addHandler(ch)
    log.info("logger is ready...")
    yield log

Here the test case setup:

from mock import Mock
from test_common import lambdaEnv


STREAM_NAME = 'stream-dev'



class TestTicketingKinesisToS3(unittest.TestCase):
    kinesis_obj = boto3.client
    event = None

    def setUp(self):
        unittest.TestCase.setUp(self)

        data1 = '{"RecordType":"Transaction","...}'

        data2 = '{"RecordType":"Transaction","...}'

        tt = [{u'eventID': u'shardId-000000000000:49574118225392902538999142319068611050512700818631688194',
               u"eventSource": u"aws:kinesis",
               u"eventVersion": u"1.0",
               u"invokeIdentityArn": u"arn:aws:iam::167856709666:role/lambda-kinesis-execution-role",
               u"eventName": u"aws:kinesis:record",
               u"eventSourceARN": u"arn:aws:kinesis:ap-southeast-2:167856709666:stream/ticketing-dev",
               u'kinesis': {
                   u'approximateArrivalTimestamp': time.time(),
                   u'data': base64.b64encode(utils.gzipDataBlob(data1)),
                   u'partitionKey': u'partitionKey-0',
                   u'sequenceNumber': u'49574118225392902538999142319068611050512700818631688194'}},
              {u'eventID': u'shardId-000000000000:49574118225392902538999142319069819976332315447806394370',
               u"eventSource": u"aws:kinesis",
               u"eventVersion": u"1.0",
               u"invokeIdentityArn": u"arn:aws:iam::167856709666:role/lambda-kinesis-execution-role",
               u"eventName": u"aws:kinesis:record",
               u"eventSourceARN": u"arn:aws:kinesis:ap-southeast-2:167856709666:stream/ticketing-dev",
               u'kinesis': {
                   u'approximateArrivalTimestamp': time.time(),
                   u'data': base64.b64encode(utils.gzipDataBlob(data2)),
                   u'partitionKey': u'partitionKey-1',
                   u'sequenceNumber': u'49574118225392902538999142319069819976332315447806394370'}},
              {u'eventID': u'shardId-000000000000:49569562892134134084378247283249679591017585731862265858',
               u"eventSource": u"aws:kinesis",
               u"eventVersion": u"1.0",
               u"invokeIdentityArn": u"arn:aws:iam::167856709666:role/lambda-kinesis-execution-role",
               u"eventName": u"aws:kinesis:record",
               u"eventSourceARN": u"arn:aws:kinesis:ap-southeast-2:167856709666:stream/ticketing-dev",
               u'kinesis': {
                   u'approximateArrivalTimestamp': time.mktime(datetime.datetime(2017, 6, 14, 16, 33, 24, 775000,
                                                                                 tzinfo=pytz.timezone(
                                                                                     'Australia/Sydney')).timetuple()),
                   u'data': base64.b64encode(
                       '\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\...'),
                   u'PartitionKey': u'mBuDKy',
                   u'sequenceNumber': u'49569562892134134084378247283249679591017585731862265858'}}
              ]

        self.event = {'Records': tt}

Now the test case:

    # @mock.patch('boto3.client', side_effect=kinesisClient)
    # def test_handler(self,mock_func):
    def test_handler(self):
        with lambdaEnv(STREAM_NAME) as log:
            import KinesisToS3
            log.debug("create context for lambda function")
            context = Mock()
            context.function_name = 'test'
            context.aws_request_id = 'd75ec968-30b4-11e8-bc5a-79557d1008ed'
            KinesisToS3.handler(self.event, context)
        # assert mock_func.called


if __name__ == '__main__':
    unittest.main()

Comments

Popular posts from this blog

How To: GitHub projects in Spring Tool Suite

Spring 3 Part 7: Spring with Databases

Parse the namespace based XML using Python