update visualization test

test/runtests.jl

@@ -30,14 +30,19 @@ end
     clean_up_temp_files()
 end
 
-@testset "2D shock detection and visualization function test" begin
+@testset "2D shock detection and 2D visualization function test" begin
     flow_data = FlowData(joinpath(DATA_DIR, "sod_shock_right_2d.tape"), false)
     point_detect_algo = ImageProcessingShockDetectionAlgo(0.5, :prewitt)
     dbscan_algo = DBSCANAlgo(0.25, 3, 10)
     fitting_algo = FittingAlgo(0.1, false)
     detection = detect(flow_data, point_detect_algo, dbscan_algo, fitting_algo)
     plot_shock_fits_over_time(flow_data, detection, true)
-    # TODO: Test heatmap function as well
+    create_heatmap_evo_with_shock(flow_data, detection, :density_field, true, false)
+    try
+        rm("density_field_evolution.gif")
+    catch e
+        @warn "Failed to clean up temporary files" exception=e
+    end
 end
 
 @testset "2D shock detection (with improved initial guess)" begin