Function to calculate energy where all parameters must be positive

Question

I have a function which requires all parameters to be positive:

def calc_energy(n, hbar=1, w=1): if hbar <= 0: raise ValueError(" hbar must be positive") if w <= 0: raise ValueError(" w must be positive") if n < 0: raise ValueError(" n must not be negative") return hbar * w * (n + 1 / 2) 

To test that this produces the correct behavior, I wrote the following test function in the same file and a main function to call it:

def test_calc_energy(): try: calc_energy(0, hbar=0) except ValueError: assert True else: assert False try: calc_energy(-1) except ValueError: assert True else: assert False try: calc_energy(0, hbar=1, w=-1) except ValueError: assert True else: assert False if __name__ == "__main__": test_calc_energy() 

Is there a better way to write the calc_energy() function or the test_calc_energy() function? Both seem rather lengthy to me.

On my laptop, I have the test in a separate file, from which I import the file with the calc_energy function, and I run pytest. One would just need to put all the code in a file, say named energy.py, then run python energy.py. It will return assertion error if the tests fail.

Answer 1

Instead of assert True you should just pass, or better, assert that the string in the exception makes sense.

Instead of assert False I typically prefer to raise AssertionError(), but that's more stylistic.

The length of the test is not particularly a problem. You can slightly reduce it by writing e.g.

try: calc_energy(0, hbar=0) raise AssertionError('Should have raised') except ValueError: pass 

However, for calc_energy, to me it smells. This is a very simple operation, and making a wild guess, those parameters might be better validated in a different location and this function reduced to one line. In a setting where this is a function internal to a module, the external interface to the module would be doing the validation, but there's no context in your question to say whether that's the case.

Is hbar the reduced Planck Constant? If so, well... it's a constant, and it doesn't need validation and should not be a function parameter.

Answer 2

In addition to Reinderien's answer, I strongly suggest relying on purpose-built test frameworks. You'll (I hope!) write more tests for other functionality in future, and manually running all .py files that have some tests executed in if __name__ == '__main__' blocks will become error-prone. You'll miss one and not notice that a test is actually failing until it surfaces somewhere else.

Test frameworks also provide additional helpers, including testing that some operation raises. With pytest your code can be greatly simplified:

# test_calc_energy.py import pytest from your_module.energy import calc_energy def test_calc_energy(): with pytest.raises(ValueError): calc_energy(0, hbar=0) with pytest.raises(ValueError): calc_energy(-1) with pytest.raises(ValueError): calc_energy(0, hbar=1, w=-1) 

nice, huh? It checks that an exception of the specified type is raised (so raising anything else or not raising at all will fail the test). You can even use with pytest.raises(ValueError, match=r" w must be positive") to check the message of the raised exception. match accepts a regex string (or an actual regex). You can run such tests simply by typing pytest in your shell - test collection is based on file and function names.

If you don't want any external dependencies (rare, but possible), there's unittest built-in module. It also provides self.assertRaises test method on testcase classes (and see assertRaisesRegex for the same message checking behavior as in pytest).

Answer 3

In addition to the suggestions of the previous answer, here are some other style suggestions.

Documentation

The PEP 8 style guide recommends adding docstrings for functions. For calc_energy, you should describe the input types and return types.

Naming

The variable names for the function inputs (n, etc.) are not very descriptive. You should use longer names or at least add documentation or comments describing what they mean in the context of your code.

Layout

The individual tests in the test_calc_energy function should be separated by blank lines to make the code easier to read and understand:

def test_calc_energy(): try: calc_energy(0, hbar=0) except ValueError: assert True else: assert False try: calc_energy(-1) except ValueError: assert True else: assert False