Improve tests and the method itself