diff --git a/bin/schedule.jl b/bin/schedule.jl
index 3b87152..01355f5 100644
--- a/bin/schedule.jl
+++ b/bin/schedule.jl
@@ -5,6 +5,8 @@ using Dagger
 using ArgParse
 using FrameworkDemo
 
+const logs_formats = ["graph", "trace", "gantt", "raw"]
+
 function parse_args(raw_args)
     s = ArgParseSettings()
 
@@ -58,9 +60,7 @@ end
 
 function (@main)(raw_args)
     args = parse_args(raw_args)
-
-    logging_required = !isnothing(args["logs-graph"]) || !isnothing(args["logs-trace"]) ||
-                       !isnothing(args["logs-gantt"]) || !isnothing(args["logs-raw"])
+    logging_required = any(x -> !isnothing(args["logs-$x"]), logs_formats)
 
     if logging_required
         FrameworkDemo.enable_logging!()
@@ -88,17 +88,11 @@ function (@main)(raw_args)
                                                           fast = fast)
     if logging_required
         logs = FrameworkDemo.fetch_logs!()
-        if !isnothing(args["logs-graph"])
-            FrameworkDemo.save_logs_graphviz(logs, args["logs-graph"])
-        end
-        if !isnothing(args["logs-trace"])
-            FrameworkDemo.save_logs_chrome_trace(logs, args["logs-trace"])
-        end
-        if !isnothing(args["logs-gantt"])
-            FrameworkDemo.save_logs_gantt(logs, args["logs-gantt"])
-        end
-        if !isnothing(args["logs-raw"])
-            FrameworkDemo.save_logs_raw(logs, args["logs-raw"])
+        for format in logs_formats
+            path = args["logs-$format"]
+            if !isnothing(path)
+                FrameworkDemo.save_logs(logs, path, Symbol(format))
+            end
         end
     end
 
diff --git a/src/visualization.jl b/src/visualization.jl
index a561bb8..8986dc6 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -4,7 +4,15 @@ import DataFrames
 import Plots
 import JSON3
 
-function save_logs_graphviz(logs, path::String)
+function save_logs(logs, ::String, ::Val{T}) where {T}
+    throw(ArgumentError("Unsupported visualization mode: `$T`"))
+end
+
+function save_logs(logs, path::String, mode::Symbol)
+    return save_logs(logs, path, Val{mode}())
+end
+
+function save_logs(logs, path::String, ::Val{:graph})
     if splitext(path)[2] == ".dot"
         open(path, "w") do io
             Dagger.show_logs(io, logs, :graphviz; color_by = :proc)
@@ -17,20 +25,20 @@ function save_logs_graphviz(logs, path::String)
     end
 end
 
-function save_logs_chrome_trace(logs, path::String)
+function save_logs(logs, path::String, ::Val{:trace})
     open(path, "w") do io
         Dagger.show_logs(io, logs, :chrome_trace)
         @info "Written logs trace to $path"
     end
 end
 
-function save_logs_gantt(logs, path::String)
+function save_logs(logs, path::String, ::Val{:gantt})
     plot = Dagger.render_logs(logs, :plots_gantt)
     Plots.savefig(plot, path)
     @info "Written logs gantt chart to $path"
 end
 
-function save_logs_raw(logs, path::String)
+function save_logs(logs, path::String, ::Val{:raw})
     if splitext(path)[2] == ".json"
         open(path, "w") do io
             JSON3.write(io, logs)
diff --git a/test/visualization.jl b/test/visualization.jl
index 0a04d0c..8144fea 100644
--- a/test/visualization.jl
+++ b/test/visualization.jl
@@ -4,19 +4,26 @@ using Graphs
 using MetaGraphs
 
 @testset "Visualization" begin
-    graph = joinpath(pkgdir(FrameworkDemo), "data/demo/datadeps/df.graphml") |>
-            FrameworkDemo.parse_graphml |> FrameworkDemo.mockup_dataflow |>
-            FrameworkDemo.get_execution_plan
+    @testset "Execution plan" begin
+        graph = joinpath(pkgdir(FrameworkDemo), "data/demo/datadeps/df.graphml") |>
+                FrameworkDemo.parse_graphml |> FrameworkDemo.mockup_dataflow |>
+                FrameworkDemo.get_execution_plan
 
-    @test nv(graph) == 7
-    @test ne(graph) == 7
-    has_edge(x::String, y::String) = Graphs.has_edge(graph, graph[x, :label],
-                                                     graph[y, :label])
-    @test has_edge("ProducerA", "TransformerAB")
-    @test has_edge("ProducerBC", "TransformerAB")
-    @test has_edge("ProducerBC", "ConsumerBC")
-    @test has_edge("ProducerBC", "TransformerC")
-    @test has_edge("ProducerBC", "ConsumerCD")
-    @test has_edge("TransformerAB", "ConsumerE")
-    @test has_edge("TransformerAB", "ConsumerCD")
+        @test nv(graph) == 7
+        @test ne(graph) == 7
+        has_edge(x::String, y::String) = Graphs.has_edge(graph, graph[x, :label],
+                                                         graph[y, :label])
+        @test has_edge("ProducerA", "TransformerAB")
+        @test has_edge("ProducerBC", "TransformerAB")
+        @test has_edge("ProducerBC", "ConsumerBC")
+        @test has_edge("ProducerBC", "TransformerC")
+        @test has_edge("ProducerBC", "ConsumerCD")
+        @test has_edge("TransformerAB", "ConsumerE")
+        @test has_edge("TransformerAB", "ConsumerCD")
+    end
+
+    @testset "Save logs" begin
+        @test_throws ArgumentError FrameworkDemo.save_logs(Dict{}(), tempname(),
+                                                           :unsupported_format)
+    end
 end