Testing Python code that produces Altair graphs

Question

Is there a recommended way to write tests for Python code that generates graphs using Altair?

A simple example is:

df = pandas.DataFrame(...)
chart = alt.Chart(df).mark_point().encode(x='...', y='...')

For simple graphs, one can just test the contents of the DataFrame, check that the Altair invocation doesn't raise an exception, and check that chart.save() also doesn't raise an exception, without testing the contents of Altair's output.

But for more complex Altair graph definitions, we really want to have our test check the output from Altair too.

What I've been doing is using a golden file (i.e. a checked-in test expectation file) for the output of chart.to_json() (the Vega JSON file).

For more complex cases, some tweaks to the JSON file are desirable, like this:

# Disabling "consolidate_datasets" makes the JSON output more
# stable, making it better to use as a golden file.  Without that
# disabled, Altair deduplicates the datasets by hash, and orders
# them in the JSON output by hash, so small changes can cause the
# datasets to be reordered.
with alt.data_transformers.enable(consolidate_datasets=False):
  json_data = chart.to_json() + '\n'
# Selector names can vary across runs of the test, so replace them
# with a placeholder.
json_data = re.sub(r'selector\d+', 'selectorXXX', json_data)
# Remove the schema URL so that the tests' golden files do not change
# across updates to Altair/Vega when only the version number in the
# URL changes.
json_data = re.sub(r'"\$schema": ".*"', '"$schema": "<removed>"', json_data)
self.assert_golden_file_equal('vega_graph.json', json_data)

Is there a better way to do this?

Answer 1

You can test specific parts of the chart output via the object oriented API:

assert chart.mark == 'bar'

Due to the fact that the Chart representation can change after displaying a chart, you likely want to use .to_dict() instead:

assert chart.to_dict()['mark'] == 'bar'

You can check for any existing chart property with this syntax, e.g.

assert chart.to_dict()['encoding']['y']['field'] == 'species'
assert chart.to_dict()['encoding']['x']['type'] == 'quantitative'
assert chart.to_dict()['encoding']['x']['aggregate'] == 'count'

If your grading setup supports sharing variables you can assign chart.to_dict() to a variable instead of recreating it each time.

You might need to check for missing keys sometimes due to Altair/Vega-Lite returning abbreviated specs when possible.