Added more tests, it's working better now