Some fixes for better Unit Testing