Add more tests

(Sorry for being late on this)