fixed failing tests
Added some tests, changed the parameter to a switch instead