code improvements + tests