From cbfc257318b3733800a6de159e4da4db78f42ebb Mon Sep 17 00:00:00 2001
From: chesedo <pieter@chesedo.me>
Date: Thu, 17 Feb 2022 18:18:29 +0200
Subject: [PATCH] feat: adding extra fields on scenario

---
 core/src/schema/scenario.rs | 60 ++++++++++++++++++++++++++++++-------
 1 file changed, 49 insertions(+), 11 deletions(-)

diff --git a/core/src/schema/scenario.rs b/core/src/schema/scenario.rs
index ce4a0e8e..db9451c6 100644
--- a/core/src/schema/scenario.rs
+++ b/core/src/schema/scenario.rs
@@ -188,24 +188,33 @@ impl Scenario {
         original: &mut ObjectContent,
         overwrites: &BTreeMap<String, Content>,
     ) -> Result<()> {
-        let original_keys: Vec<_> = original.fields.keys().collect();
+        let original_keys: Vec<String> = original
+            .fields
+            .keys()
+            .into_iter()
+            .map(ToOwned::to_owned)
+            .collect();
         let original_len = original_keys.len();
 
-        for overwrite in overwrites.keys() {
-            if !original_keys.contains(&overwrite) {
-                return Err(anyhow!(
-                    "'{}' is not a field on the object, therefore it cannot be included",
-                    overwrite
-                ));
+        for (overwrite_name, overwrite_content) in overwrites {
+            if !original_keys.contains(&overwrite_name) {
+                // Cannot add include only fields
+                if let Content::Empty(_) = overwrite_content {
+                    return Err(anyhow!(
+                        "'{}' is not a field on the object, therefore it cannot be included",
+                        overwrite_name
+                    ));
+                }
+
+                original
+                    .fields
+                    .insert(overwrite_name.to_string(), overwrite_content.clone());
             }
         }
         let (keep_fields, trim_fields): (Vec<_>, Vec<_>) = original_keys
             .into_iter()
             .partition(|c| overwrites.contains_key(c.as_str()));
 
-        let trim_fields: Vec<_> = trim_fields.into_iter().map(ToOwned::to_owned).collect();
-        let keep_fields: Vec<_> = keep_fields.into_iter().map(ToOwned::to_owned).collect();
-
         for trim_field in trim_fields {
             debug!("removing field '{}'", trim_field);
             original.fields.remove(trim_field.as_str());
@@ -391,7 +400,7 @@ mod tests {
 
     #[test]
     #[should_panic(expected = "'null' is not a field on the object")]
-    fn build_filter_extra_field() {
+    fn build_filter_extra_field_include() {
         let scenario = Scenario {
             namespace: namespace!({
                 "collection1": {
@@ -408,6 +417,35 @@ mod tests {
         scenario.build().unwrap();
     }
 
+    #[test]
+    fn build_filter_extra_field() {
+        let scenario = Scenario {
+            namespace: namespace!({
+                "collection1": {
+                    "type": "object",
+                    "nully": {"type": "null"},
+                    "stringy": {"type": "string", "pattern": "test"}
+                },
+                "collection2": {}
+            }),
+            scenario: scenario!({"collection1": {"number": {"type": "number", "constant": 4}}}),
+            name: "test".to_string(),
+        };
+
+        let actual = scenario.build().unwrap();
+        let expected = namespace!({
+            "collection1": {
+                "type": "object",
+                "number": {
+                    "type": "number",
+                    "constant": 4
+                }
+            }
+        });
+
+        assert_eq!(actual, expected);
+    }
+
     #[test]
     #[should_panic(expected = "cannot select fields to include from a non-object")]
     fn build_filter_field_scalar() {