24 Nisan 2022 Pazar

How to mock socket object used with "with" keyword on python Unit Tests

In a development process, you have to make sure that all your working code have to remain as working as intended to be. Just like when you first wrote them. So maybe you think that you have to write unit test for your functions, classes, etc. But most of the time you have to isolate your application from external and variable inputs. In a situation like this, you will basically use mocking mechanisms to set inputs that outside from the units (outside from a function or class).

So let's start with basic mocking example and see how it goes.

We simply mocked randint function with patch decorator and change its return value. But the things are different for a socket object used with "with" keyword. So how can be different? First of all we have to look at context managers and how we can create a class usable with "with" keyword. For code simplicity's sake, example includes "open" built-in function which already supports this behaviour.

The magic lays underneath __enter__ and __exit__ functions. These are magic functions that you can implement for operators and built-in keywords like "with". When we use "with" keyword and create an instance from a class, __enter__ magic function is also executed and returned an object, after that it was assigned to the variable that is declarated after "as" keyword. When the code under "with" keyword is finished, __exit__ function executed before the following code. So how can we write a unit test for this?

At unit test side, we mocked the File object with MagicMock object which is a Mock object that already has most of the default magic method implementations. After that we changed return value of readline which is a function of object which returned from __enter__ function of File instance. As a result, we changed file input and tested the code under "else" statement.

All these information above, we can write a unit test for socket object, now.

At the end, like the File example above, we mocked recv function and tested the code under "else" statement.