Fix: make tests nicer