diff --git a/src/ast/dml/insert.rs b/src/ast/dml/insert.rs
index 74ebad5..26a05da 100644
--- a/src/ast/dml/insert.rs
+++ b/src/ast/dml/insert.rs
@@ -80,14 +80,14 @@ mod tests {
             .build();
 
         assert_eq!(
-            insert_query,
-            InsertQuery {
+            SQLStatement::from(insert_query),
+            SQLStatement::DML(DMLStatement::InsertQuery(InsertQuery {
                 into_table: Some(TableName::new(None, "table".into())),
                 columns: vec!["a".into(), "b".into()],
                 data: InsertData::Values(vec![InsertValue {
                     list: vec![Some(SQLExpression::String("a".into()))],
                 }]),
-            }
+            }))
         );
     }
 }
diff --git a/src/ast/dml/select.rs b/src/ast/dml/select.rs
index 15f4121..06accff 100644
--- a/src/ast/dml/select.rs
+++ b/src/ast/dml/select.rs
@@ -299,8 +299,8 @@ mod tests {
             .build();
 
         assert_eq!(
-            select_query,
-            SelectQuery {
+            SQLStatement::from(select_query),
+            SQLStatement::DML(DMLStatement::SelectQuery(SelectQuery {
                 select_items: vec![SelectKind::SelectItem(SelectItem {
                     item: Some(SQLExpression::String("a".into())),
                     alias: None
@@ -314,7 +314,7 @@ mod tests {
                 limit: None,
                 offset: None,
                 has_aggregate: false,
-            }
+            }))
         );
     }
 }