How to mock a context manager and replace one or more of its methods

Question

I am trying to mock a class which is used as a context manager that makes network calls. Specifically, it's Read method returns the value of a PLC's tag over the network.

Here is an example of the function I am trying to test in main.py:

from pylogix import PLC

def read_counter(counter_entry):
    with PLC() as comm:
        count = comm.Read(counter_entry['tag'])
        if count.Status == 'Success':
            log_count(count.Value)
            counter_entry['last_count'] == count.Value

My test looks like this:

import unittest
from unittest.mock import Mock, patch, MagicMock
from random import randint

import main

class MockComm():

    def __init__(self):
        self.tag_dict = {}

    def __enter__(self):
        print('__enter__ called')
        return self

    def __exit__(self):
        print('__exit__ called')

    def add(self, tag, value):
        print('Added ', tag, ' as ', value)
        self.tag_dict[tag] = value

    def Read(self, tag):
        print('Reading ', tag)
        if tag in self.tag_dict:
            print('returned ', self.tag_dict[tag])
            return Response(tag, Value=self.tag_dict[tag])
        else:
            print('returned Connection Failure')
            return Response(tag, Value=None, Status='Connection failure')


class Response():
    def __init__(self, tag, Value=None, Status='Success'):
        self.TagName = tag
        self.Value = Value
        self.Status = Status

class ReadPylogixCounterTestSuit(unittest.TestCase):

    def setUp(self):
        self.counter_entry = {
            'type': 'pylogix_counter',
            'tag': 'Program:Production.DailyTotal',
            'Part_Type_Tag': 'Stn010.PartType',
            'Part_Type_Map': {'0': '50-4865', '1': '50-5081'},
            # used internally to track the readings
            'nextread': 0,      # timestamp of the next reading
            'lastcount': 0,     # last counter value
            'lastread': 0       # timestamp of the last read
        }

    @patch('main.PLC', new_callable=MockComm)
    def test_read_tag(self, mock_PLC):

        FIRST_COUNTER_VALUE = randint(1, 2500)
        mock_PLC.add(self.counter_entry['tag'], FIRST_COUNTER_VALUE)
        mock_PLC.add(self.counter_entry['Part_Type_Tag'], 0)

        main.read_counter(self.counter_entry)

        assert self.counter_entry['lastcount'] == FIRST_COUNTER_VALUE

EDIT: Update the code in the question to be a complete runnable example.
When I step through the above code, mock_PLC is an instance of Mock_Comm. As soon as I hit the with PLC() as comm line I get the following traceback:

=====================================================================
ERROR: test_read_tag (test_read_counter.ReadPylogixCounterTestSuit)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "C:\Users\cstrutto\AppData\Local\Programs\Python\Python37\lib\unittest\mock.py", line 1256, in patched
    return func(*args, **keywargs)
  File "c:\programing\python-testing\test_context_manager\test_read_counter.py", line 81, in test_read_tag
    main.read_counter(self.counter_entry)
  File "c:\programing\python-testing\test_context_manager\main.py", line 4, in read_counter
    with PLC() as comm:
TypeError: 'MockComm' object is not callable

----------------------------------------------------------------------
Ran 1 test in 34.908s

FAILED (errors=1)

I have the required __enter__ and __exit__ methods. Why is it not working as a context manager?

May as well `return self`. You'll need to define an `__exit__` too. — wim, Aug 31 '20 at 21:31
Yes, the `__enter__` method is that simple. I have updated the question to have a fully testable example. — cstrutton, Sep 01 '20 at 17:36

score 0 · Answer 1 · answered Sep 01 '20 at 20:18

I have not been able to replace the object like I was trying to do above however, I did find a way to set the side_effect for the Read method in the mock object.

I will still try to replace the object with my own, but for now, the following code works:

@patch('main.PLC', autospec=True)
def test_read_tag(self, mock_PLC):

    FIRST_COUNTER_VALUE = randint(1, 2500)

    def mock_Read(tag):
        
        tag_dict={
            self.counter_entry['tag']: FIRST_COUNTER_VALUE,
            self.counter_entry['Part_Type_Tag']: 0
        }

        print('Reading ', tag)
        if tag in tag_dict:
            print('returned ', tag_dict[tag])
            return Response(tag, Value=tag_dict[tag])
        else:
            print('returned Connection Failure')
            return Response(tag, Value=None, Status='Connection failure')

    mock_client = MagicMock(spec=main.PLC)
    mock_client.Read.side_effect = mock_Read
    mock_client.__enter__.return_value = mock_client
    mock_PLC.return_value = mock_client

    main.read_counter(self.counter_entry)

How to mock a context manager and replace one or more of its methods

1 Answers1