[BUGFIX] Get unit tests working again