Tests related code update