Fixing tests and adding helper functions