From 53a8b6aa977bc6c051625a812009184f78da597d Mon Sep 17 00:00:00 2001
From: Muriithi Frederick Muriuki
Date: Wed, 4 Aug 2021 10:05:33 +0300
Subject: 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
---
 gn3/db/traits.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

(limited to 'gn3')

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.
-- 
cgit 1.4.1