Restructuring and adding a test