Simplify switch statement
Fix error in tests
Rework tests