From 3d38d6e8c5f39c4da6011876b077efbb97b1bad0 Mon Sep 17 00:00:00 2001
From: Gavin Peacock <gpeacock@adobe.com>
Date: Wed, 17 Jul 2024 15:32:47 -0700
Subject: [PATCH] Ensure Ingredient data_types make it to the store and back.

---
 sdk/src/ingredient.rs |  2 ++
 sdk/src/manifest.rs   | 42 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/sdk/src/ingredient.rs b/sdk/src/ingredient.rs
index 90f2e1c45..41d0f148d 100644
--- a/sdk/src/ingredient.rs
+++ b/sdk/src/ingredient.rs
@@ -1082,6 +1082,7 @@ impl Ingredient {
         ingredient.metadata = ingredient_assertion.metadata;
         ingredient.description = ingredient_assertion.description;
         ingredient.informational_uri = ingredient_assertion.informational_uri;
+        ingredient.data_types = ingredient_assertion.data_types;
         Ok(ingredient)
     }
 
@@ -1265,6 +1266,7 @@ impl Ingredient {
         ingredient_assertion
             .informational_uri
             .clone_from(&self.informational_uri);
+        ingredient_assertion.data_types.clone_from(&self.data_types);
         claim.add_assertion(&ingredient_assertion)
     }
 
diff --git a/sdk/src/manifest.rs b/sdk/src/manifest.rs
index 1a7d049de..57f0792ec 100644
--- a/sdk/src/manifest.rs
+++ b/sdk/src/manifest.rs
@@ -2288,14 +2288,24 @@ pub(crate) mod tests {
             "format": "text/plain",
             "relationship": "inputTo",
             "data": {
-              "format": "text/plain",
-              "identifier": "prompt.txt",
-              "data_types": [
+                "format": "text/plain",
+                "identifier": "prompt.txt",
+                "data_types": [
+                    {
+                    "type": "c2pa.types.generator.prompt"
+                    }
+                ]
+            }
+        },
+        {
+            "title": "Custom AI Model",
+            "format": "application/octet-stream",
+            "relationship": "inputTo",
+            "data_types": [
                 {
-                  "type": "c2pa.types.generator.prompt"
+                    "type": "c2pa.types.model"
                 }
-              ]
-            }
+            ]
           }
         ]
     }"#;
@@ -2347,7 +2357,8 @@ pub(crate) mod tests {
         let (format, image) = m.thumbnail().unwrap();
         assert_eq!(format, "image/jpeg");
         assert_eq!(image.to_vec(), b"my value");
-        assert_eq!(m.ingredients().len(), 2);
+        assert_eq!(m.ingredients().len(), 3);
+        // Validate a prompt ingredient (with data field)
         assert_eq!(m.ingredients()[1].relationship(), &Relationship::InputTo);
         assert!(m.ingredients()[1].data_ref().is_some());
         assert_eq!(m.ingredients()[1].data_ref().unwrap().format, "text/plain");
@@ -2356,6 +2367,14 @@ pub(crate) mod tests {
             m.ingredients()[1].resources().get(id).unwrap().into_owned(),
             b"pirate with bird on shoulder"
         );
+        // Validate a custom AI model ingredient.
+        assert_eq!(m.ingredients()[2].title(), "Custom AI Model");
+        assert_eq!(m.ingredients()[2].relationship(), &Relationship::InputTo);
+        assert_eq!(
+            m.ingredients()[2].data_types().unwrap()[0].asset_type,
+            "c2pa.types.model"
+        );
+
         // println!("{manifest_store}");
     }
 
@@ -2400,7 +2419,7 @@ pub(crate) mod tests {
         let (format, image) = m.thumbnail().unwrap();
         assert_eq!(format, "image/jpeg");
         assert_eq!(image.to_vec(), b"my value");
-        assert_eq!(m.ingredients().len(), 2);
+        assert_eq!(m.ingredients().len(), 3);
         assert_eq!(m.ingredients()[1].relationship(), &Relationship::InputTo);
         assert!(m.ingredients()[1].data_ref().is_some());
         assert_eq!(m.ingredients()[1].data_ref().unwrap().format, "text/plain");
@@ -2409,6 +2428,13 @@ pub(crate) mod tests {
             m.ingredients()[1].resources().get(id).unwrap().into_owned(),
             b"pirate with bird on shoulder"
         );
+        // Validate a custom AI model ingredient.
+        assert_eq!(m.ingredients()[2].title(), "Custom AI Model");
+        assert_eq!(m.ingredients()[2].relationship(), &Relationship::InputTo);
+        assert_eq!(
+            m.ingredients()[2].data_types().unwrap()[0].asset_type,
+            "c2pa.types.model"
+        );
         // println!("{manifest_store}");
     }