test improvements
simplified code
Improved code, make it work