Writing Unit Tests in Python with Pytest

Vedant Nibandhe
6 min readNov 10, 2020

Before writing this article, I assume you have some experience with programming. To be more specific, Python. Well, if you don’t you can always come back later.

If you ask me, personally, tutorials would be a weird way to learn a programming language, there are possibilities of getting stuck in the tutorial-loop for a long time, and it's pretty frustrating. Now, I am not saying you should not follow tutorials. Just understand the basics like syntax, language-specific rules, and a handful of problems from the tutorials, and you are set to start coding. Practicing a plethora of problems is the way to go. HackerRank is a great place to start solving and my personal favorite. If you get stuck, you can always google your query out.

Anyway, I digressed. Coming back to HackerRank, I want to address those who already use this platform. I suppose this must have happened to you:

  1. You see an exciting problem.
  2. You feel optimistic about it.
  3. You conceptualize the logic of how you are going to solve it using your preferred programming language.
  4. You program it.
  5. You get the dopamine high and press the ‘Submit’ button. Anddd boom!!

You see a ton of test cases failing. Now what? you don’t feel so good. You figure something out, alter your code a bit, re-submit it. Still, some of the test-cases failing(probably corner cases) and you have no idea what’s failing. You doubt if you are ever going to be the decorated programmer you wished to be. Don’t worry. It happens to the best of programmers and will happen to you too in 8/10 cases.

Personally, when I started the journey, I would feel handling corner cases(these are just inputs which have the potential to break your code, more on this later!) in your code is just useless. I mean, why would the input to my code change? Right? Wrong. In a real environment(i.e app, website, etc) everything is fragile. You need to imagine every possible input that can enter your code. And that’s why a good programmer usually writes tests first, and code later. Initially, all or a few tests may not pass, but eventually, the goal is to write your code in such a way so that it passes most of the test cases. This is called Test-Driven Development(TDD).

Unit tests are tests you write for the smallest testable component of your program(i.e a method, function). But why unit test?

Yup, you gotta

Unit tests improve the quality of the code. It identifies every bug that may come up before code is sent further for integration testing.

Let’s jump into it and see how you can write tests in python. We are going to use a module named ‘pytest’ which is not a part of Python’s standard library. We can install using the pip command on Windows.

pip install pytest

For Linux and MacOS, you can install it using pip3.
To keep it simple, let’s write unit tests for something like Dividing 2 integers to return an integer.

It may sound pretty basic, but let’s focus on the gravity of the situation with this. Let’s visualize the code for dividing two integers.
Before writing unit tests with ‘pytest’, let's note somethings followed while writing a unit test.

  1. Say your python file name is ‘divide_integers.py’. So, your test file name would be ‘test_divide_integers.py’.
  2. Each test case in ‘test_divide_integers.py’ would be a function testing for different values of input with the desired output.
  3. Each function in your ‘test_divide_integers.py’ would start with the prefix ‘test_….()’.
  4. Because you have a file with the prefix test, pytest knows it will contain unit tests for it to run. The same principles apply to the class and method names inside the file.

Now that we know quite some things about pytest, let's get into the implementation. Assuming you have installed pytest on your system, here the following steps:

  1. Create an empty folder(preferably with the name ‘test’ for best practices).
  2. Create a file called ‘test_divide_integers.py’, which is the file we’ll write our tests in.
  3. Create a file called ‘divide_integers.py’, which is the functionality we want to test.

Note: pytest uses the ‘assert’ statement for equality checks. Below, 5 is the desired output that we want when we pass 10 and 2 to the function we are testing.

All Set!

Now, let’s write our first unit test.

#test_divide_integers.pyfrom divide_integers import dividedef test_case1():
assert 5 == divide(10,2) #calls the function we want to test
#divide_integers.pydef divide(numerator, denominator):
return numerator/denominator

Just cmd into the folder where you have saved all these files. You can run your unit test(Just one test case for now) with the following commands with ‘py.test test_filename.py -v’:

py.test test_divide_integers.py -v
or
python -m pytest test_divide_integers.py

You would see something like this on your terminal. You can append the ‘- -disable-warnings’ flag to eliminate warnings.

Great! We just wrote our first unit test. Wait. There are still some things that can break the code. Let’s discuss them.

  1. What if you get 0 passed as a denominator to your function?
  2. What if either your numerator or your denominator is None type?
  3. What if you want your output as an integer and your numerator is not divisible by your denominator?
  4. Since we are talking about the division of two integers, what if you get input as a floating point number(this is not a test-case, this has to do more with types but let's just call it a test-case for now).

All these are corner cases that I talked about at the beginning of the article.

Let’s design test cases(add new test cases to test_divide_integers.py) now and let's run our program to see if it works well. The new test cases would be individual functions written in the following way:

#test_divide_integers.pyfrom divide_integers import dividedef test_case1():
assert 5 == divide(10,2)
def test_case2():
assert 'Not defined' == divide(10,0)
def test_case3():
assert 3 == divide(10,3)
def test_case4():
assert 'Numerator/denominator value is None' == divide(None, 3)

Disappointing, right? Our program passes on only one standard test case that’s it. ‘pytest’ being the elegant library it is, gives us where our code went wrong. The comparison between the desired output and the actual value is given as well(test_case3 in FAILURES section). Without further ado, let’s jump in and change our code to handle all the test cases.

#divide_integers.py
def divide(numerator, denominator):
if denominator == 0: #test_case2
return 'Not defined'
elif denominator is None or numerator is None: #test_case4
return 'Numerator/denominator value is None'
return int(numerator/denominator) #test_case3 and test_case1

On running the tests again, we get the following output:

Kudos!!! You just learned how to write unit tests in Python. If you guys want to try, you can experiment by expanding the problem to multiplication as well. You can also use the ‘pytest’ to write unit tests for whichever algorithm you guys are solving, the structure of ‘pytest’ remains the same. You can hop onto https://docs.pytest.org/en/latest/ for ‘pytest’ documentation if you want to dive deeper. You can email me at vnibandhe@gmail.com to ask stuff about unit testing and ML related stuff. See you guys!

--

--

Vedant Nibandhe

Software Engineer. Interested in sports, tech, and business.