Improving unit tests
Minor corrections