Add tests for post-processing functions

Issues:
https://github.com/genenetwork/gn-gemtext-threads/blob/main/topics/gn1-migration-to-gn2/clustering.gmi

* Add missing tests for some post-processing functions

1 files changed, 8 insertions, 5 deletions

diff --git a/gn3/db/traits.py b/gn3/db/traits.py
index 902eb8b..ce6298f 100644
--- a/gn3/db/traits.py
+++ b/gn3/db/traits.py
@@ -160,11 +160,14 @@ def set_confidential_field(trait_info):
     """Post processing function for 'Publish' trait types.
 
     It sets the value for the 'confidential' key."""
-    return {
-        **trait_info,
-        "confidential": 1 if (
-            trait_info.get("pre_publication_description", None)
-            and not trait_info.get("pubmed_id", None)) else 0}
+    if trait_info["type"] == "Publish":
+        return {
+            **trait_info,
+            "confidential": 1 if (
+                trait_info.get("pre_publication_description", None)
+                and not trait_info.get("pubmed_id", None)) else 0}
+    else:
+        return trait_info
 
 def retrieve_probeset_trait_info(trait_data_source: Dict[str, Any], conn: Any):
     """Retrieve trait information for type `ProbeSet` traits.