Few more changes, test actually works...