Refactor tests to use new visualisation factory with correct data format