Tests + bug fix