Added tests, optimized logic

Hey, I realized that I forgot to include the test files that I used
while getting this to work. I also realized I could make some general
improvements on the logic.  Went ahead and did it.