cassandra/diagram: Fix assertion to keep mypy happy

The assertion was failing when executing the code so use a different
one to keep mypy happy.

Also add test to ensure the diagram generation code is covered.