Fixing Tests and improving code.