Updated to have more tests and improving existing