Added the second test and reorganized the code a bit.