Fixing Junit tests - note this version is NOT production ready