Refactor tests with builder class