Tests refactoring and improvements