Add tests for `visualise.py`

[jfkcooper]

We may not be able to get 100% coverage here since we don't necessarily want to write out to file.