Fixed bug with segue
Just to save current state
Added some working unit tests