Unit testing tweak