Search code examples
pythonunit-testingmagicmock

Mock entire python class


I'm trying to make a simple test in python, but I'm not able to figure it out how to accomplish the mocking process.

This is the class and def code:

class FileRemoveOp(...)
    @apply_defaults
    def __init__(
            self,
            source_conn_keys,
            source_conn_id='conn_default',
            *args, **kwargs):
        super(v4FileRemoveOperator, self).__init__(*args, **kwargs)
        self.source_conn_keys = source_conn_keys
        self.source_conn_id = source_conn_id


    def execute (self, context)
          source_conn = Connection(conn_id)
          try:
              for source_conn_key in self.source_keys:
                  if not source_conn.check_for_key(source_conn_key):    
                      logging.info("The source key does not exist")  
                  source_conn.remove_file(source_conn_key,'')
          finally:
              logging.info("Remove operation successful.")

And this is my test for the execute function:

@mock.patch('main.Connection')
def test_remove_execute(self,MockConn):
    mock_coon = MockConn.return_value
    mock_coon.value = #I'm not sure what to put here#
    remove_operator = FileRemoveOp(...)
    remove_operator.execute(self)

Since the execute method try to make a connection, I need to mock that, I don't want to make a real connection, just return something mock. How can I make that? I'm used to do testing in Java but I never did on python..


Solution

  • First it is very important to understand that you always need to Mock where it the thing you are trying to mock out is used as stated in the unittest.mock documentation.

    The basic principle is that you patch where an object is looked up, which is not necessarily the same place as where it is defined.

    Next what you would need to do is to return a MagicMock instance as return_value of the patched object. So to summarize this you would need to use the following sequence.

    • Patch Object
    • prepare MagicMock to be used
    • return the MagicMock we've just created as return_value

    Here a quick example of a project.

    connection.py (Class we would like to Mock)

    class Connection(object):                                                        
        def execute(self):                                                           
            return "Connection to server made"
    

    file.py (Where the Class is used)

    from project.connection import Connection                                        
    
    
    class FileRemoveOp(object):                                                      
        def __init__(self, foo):                                                     
            self.foo = foo                                                           
    
        def execute(self):                                                           
            conn = Connection()                                                      
            result = conn.execute()                                                  
            return result    
    

    tests/test_file.py

    import unittest                                                                  
    from unittest.mock import patch, MagicMock                                       
    from project.file import FileRemoveOp                                            
    
    class TestFileRemoveOp(unittest.TestCase):                                       
        def setUp(self):                                                             
            self.fileremoveop = FileRemoveOp('foobar')                               
    
        @patch('project.file.Connection')                                            
        def test_execute(self, connection_mock):
            # Create a new MagickMock instance which will be the
            # `return_value` of our patched object                                     
            connection_instance = MagicMock()                                        
            connection_instance.execute.return_value = "testing"
    
            # Return the above created `connection_instance`                     
            connection_mock.return_value = connection_instance                       
    
            result = self.fileremoveop.execute()                                     
            expected = "testing"                                                     
            self.assertEqual(result, expected)                                       
    
        def test_not_mocked(self):
            # No mocking involved will execute the `Connection.execute` method                                                   
            result = self.fileremoveop.execute()                                     
            expected = "Connection to server made"                                   
            self.assertEqual(result, expected)