Code refactored and further tests added
Further tidying up