Added more tests and flow control. Not yet tested yet, though!