added test and simplified a bit

Added a simple test and (hopefully) simplified the code a bit. Nothing too fancy :D