Updated codes to pass the unit tests