[test] Add more unit tests!

note: Branch is rebased as fixed are made to master