WIP: improved functionality for testing