Use more scala idioms
Add unit tests
Make tests more structured