Closing #6 with tests