More unit tests an some misc improvements