2 files changed, 8 insertions, 7 deletions
diff --git a/gn3/computations/datasets.py b/gn3/computations/datasets.py
index b405e55..3664e4a 100644
--- a/gn3/computations/datasets.py
+++ b/gn3/computations/datasets.py
@@ -1,6 +1,5 @@
 """module contains the code all related to datasets"""
 import json
-from unittest import mock
 from math import ceil
 from collections import defaultdict
 
@@ -16,7 +15,8 @@ from gn3.settings import GN2_BASE_URL
 
 def retrieve_trait_sample_data(dataset,
                                trait_name: str,
-                               group_species_id=None,) -> List:
+                               database,
+                               group_species_id=None) -> List:
     """given the dataset id and trait_name fetch the\
     sample_name,value from the dataset"""
 
@@ -35,7 +35,8 @@ def retrieve_trait_sample_data(dataset,
 
     if dataset_query:
         formatted_query = dataset_query % sample_query_values[dataset_type]
-        results = fetch_from_db_sample_data(formatted_query, mock.Mock())
+
+        results = fetch_from_db_sample_data(formatted_query, database)
 
     return results
 
diff --git a/gn3/computations/traits.py b/gn3/computations/traits.py
index bbb3073..38072ed 100644
--- a/gn3/computations/traits.py
+++ b/gn3/computations/traits.py
@@ -7,7 +7,7 @@ def compute_sum(rhs_val: int, lhs_val: int) -> int:
     return rhs_val + lhs_val
 
 
-def fetch_trait(dataset, trait_name: str) -> dict:
+def fetch_trait(dataset, trait_name: str, database) -> dict:
     """this method creates a trait by\
     fetching required data given the\
     dataset and trait_name"""
@@ -17,20 +17,20 @@ def fetch_trait(dataset, trait_name: str) -> dict:
         "trait_name": trait_name
     }
 
-    trait_data = get_trait_sample_data(dataset, trait_name)
+    trait_data = get_trait_sample_data(dataset, trait_name, database)
 
     created_trait["trait_data"] = trait_data
 
     return created_trait
 
 
-def get_trait_sample_data(trait_dataset, trait_name) -> dict:
+def get_trait_sample_data(trait_dataset, trait_name, database) -> dict:
     """first try to fetch the traits sample data from redis if that\
     try to fetch from the traits dataset redis is only  used for\
     temp dataset type which is not used in this case """
 
     sample_results = retrieve_trait_sample_data(
-        trait_dataset.id, trait_dataset.type, trait_name)
+        trait_dataset, trait_name, database)
 
     trait_data = {}