In-module tests with imports

There is a nice pattern for a scripts: you write their tests right in the script. It’s allow to bypass any import issues and just ship code and tests together.

This approach works perfectly if you have only simple tests:

def my_production_function(x):
  if x > 0:
    return x+1
  else:
    raise Exception("Negative")


### TESTS HERE


def test_1():
  assert my_production_function(1) == 2

But what we can’t import anything test-related in the top-level of the module. May be this code will run inside an environment where we don’t have pytest. With simple tests it’s simple: you just have unused function (test_1, etc), which are ignored by normal code execution. If you write import pytest at the top of the file, you must have pytest in production, which is odd.

But, let’s say you want to test your function for exceptions. How? My old approach was to import pytest inside test:

def test_0():
   import pytest
   pytest.raises(Exception):
      my_production_function(0)

It was okay… until it’s a one or two tests. If you have more of them, multiple imports slows everything down, and are ugly.

Today I found beautiful solution. Look at this!

def setup_module():
  global pytest
  import pytest


def test_0():
  with pytest.raises(Exception):
    my_production_function(0)


def test_negative():
  with pytest.raises(Exception):
    my_production_function(-1)

The trick is in setup_module where you can have global (top-level) variables. And module is a variable! Pytest runs setup_module before running tests, so, it’s importing pytest module. It’s a neat pytest-style trick.