Unittest - How to test for sys.exit

How I have managed to test a function that calls sys.exit(1)

Table of Contents
  1. The Test
  2. The code

In opsdroid, we were trying to update pyyaml to version 4.2b1 in order to fix the security vulnerability of version 3.12 that allows users to run python code from within a .yaml file. The fix was rather easy, we simply had to replace yaml.loader(stream) to yaml.loader(stream, Loader=SafeLoader) but I wanted to add a test that shows that this fix does work.

One of the first things we do in opsdroid is to load the file config.yaml to get the configuration for the bot. Before the update, you could run python code from within the config.yaml like this test: !!python/object/apply:os.system ["echo 'Oops!';"] - this will print Oops! into the shell

After the update and the fix, when a user tries to run python code from within a yaml file the following happens:

1could not determine a constructor for the tag ',2002:python/object/apply:os.system'
2in "/Users/<user>/Library/Application Support/opsdroid/configuration.yaml"
3#sys.exit(1) called

Since the next version of pyyaml might change how things work, I wanted to create a test to check if everything would work as expected and no python code could be run from within a yaml file.

There was only one problem.

Unittest calls sys.exit() when all the tests finish running, so the test that I created was causing all sort of issues with the rest of the tests - some of them were passing but most of them failed.

After a bit of research on StackOverflow and a bit of trial and error I came up with a way to test the fix and make the rest of the tests work properly.

The Test

I came across this post on tutorials point - How do you test that a Python function throws an exception, Manogna suggested to add unittest.main(exit=False) to the test and that solved my issues.

So the test ended up looking like this:

1def test_load_exploit(self):
2 opsdroid, loader = self.setup()
3 with self.assertRaises(SystemExit):
4 config = loader.load_config_file(
5 [os.path.abspath("tests/configs/include_exploit.yaml")])
6 self.assertLogs('_LOGGER', 'critical')
7 self.assertRaises(YAMLError)
8 unittest.main(exit=False)

Basically, this test does five things:

  1. loads a yaml file that contains !!python/object/apply:os.system ["echo 'Oops!';"]
  2. asserts if our Logger will log a critical message
  3. asserts if yaml.YAMLError is raised
  4. asserts if SystemExit is raised - due to sys.exit() call
  5. prevents unittest from exiting

The code

The code that these test covers is shown below, I decided to add it here just in case you want to know exactly what the code does and hopefully, it can help you somehow.

2 with open(config_path, 'r') as stream:
3"Loaded config from %s."), config_path)
4 return yaml.load(stream, Loader=yaml.SafeLoader)
5except yaml.YAMLError as error:
6 _LOGGER.critical(error)
7 sys.exit(1)
8except FileNotFoundError as error:
9 _LOGGER.critical(error)
10 sys.exit(1)

As you can see, we try to open the yaml file and load it using yaml.load(stream, Loader=yaml.SafeLoader), but if an exception is raised we just log the error and call sys.exit since the bot can't operate without any configuration.

Finally, if you would like to check the whole project, have a look at opsdroid on GitHub we are always looking for new contributors no matter the experience!


0 Like 0 Comment

You might also like these

An example from opsdroid on how to test if a logging call was made successfully.

Read More

Test: Was Logging called

Test: Was Logging called

This is an example on how to use the side_effect function from the unittest module to test if the aiohttp exception ClientOSError was raised.

Read More

Test for aiohttp's ClientErrorOS

Test for aiohttp's ClientErrorOS

While working on adding tests to Pyscript I came across a use case where I had to check if an example image is always generated the same.

Read More

How to compare two images using NumPy

How to compare two images using NumPy

Learn what additional permissions you need to add to your user to get django to run tests with a postgresql database.

Read More

Fix django postgresql permissions denied on tests

Fix django postgresql permissions denied on tests