Fix: refactor tests