How I have managed to test a function that calls sys.exit(1)
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, 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 'tag:yaml.org,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.
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.
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:
- loads a yaml file that contains
!!python/object/apply:os.system ["echo 'Oops!';"]
- asserts if our Logger will log a critical message
- asserts if
- asserts if
SystemExitis raised - due to sys.exit() call
- prevents unittest from exiting
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.
1try:2 with open(config_path, 'r') as stream:3 _LOGGER.info(_("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!