diff --git a/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py b/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py
index 384a82de0..b21bee8ad 100644
--- a/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py
+++ b/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py
@@ -194,6 +194,8 @@ def genrete_treat_subgraphs(self, response: ARAXResponse, top_drugs: pd.DataFram
                     node_info = xdtdmapping.get_node_info(node_id=drug_canonical_id)
                 except:
                     continue
+                if not node_info:
+                    continue
                 drug_categories = [node_info.category]
                 # add the node to the knowledge graph
                 drug_name = node_info.name
diff --git a/code/ARAX/test/test_ARAX_expand.py b/code/ARAX/test/test_ARAX_expand.py
index 528b09945..38563a98d 100644
--- a/code/ARAX/test/test_ARAX_expand.py
+++ b/code/ARAX/test/test_ARAX_expand.py
@@ -1302,7 +1302,7 @@ def test_xdtd_curie_not_in_db():
             "t_edge": {
                 "object": "disease",
                 "subject": "chemical",
-                "predicates": ["biolink:affects"],
+                "predicates": ["biolink:treats"],
                 "knowledge_type": "inferred"
             }
         }