added tests and fixed class a bit