diff --git a/src/cmd/query.rs b/src/cmd/query.rs
index b21d5ff..ff789e8 100644
--- a/src/cmd/query.rs
+++ b/src/cmd/query.rs
@@ -27,7 +27,7 @@ pub fn query(matches: &ArgMatches) -> Result<(), anyhow::Error> {
                     println!("Display query result: {}", df_query);
                 }
 
-                if matches.contains_id(arg::SAVE) {
+                if matches.get_flag(arg::SAVE) {
                     let save_dir = Path::new(arg::value(arg::SAVE_DIR, matches)?);
 
                     // If argument 'FILE_TYPE' is not present the default value 'csv' will be used
diff --git a/tests/cmd/test_query.rs b/tests/cmd/test_query.rs
index 149bab2..266de16 100644
--- a/tests/cmd/test_query.rs
+++ b/tests/cmd/test_query.rs
@@ -80,6 +80,7 @@ fn test_query_save_png() {
     cmd.args([
         "query",
         test_query,
+        "--display",
         "--save",
         "--save-dir",
         save_dir,