Adding in code to get tests to run