Unit Tests && simple fixes