Add assert statement to check model structure on model_visualization_…

…test (#20208)

* Add assert statement to check model structure

* Apply pytest on model_visualization_test