about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--gn3/computations/traits.py4
-rw-r--r--tests/unit/computations/test_trait.py5
2 files changed, 6 insertions, 3 deletions
diff --git a/gn3/computations/traits.py b/gn3/computations/traits.py
index 7386745..bbb3073 100644
--- a/gn3/computations/traits.py
+++ b/gn3/computations/traits.py
@@ -1,4 +1,5 @@
 """module contains all operating related to traits"""
+from gn3.computations.datasets import retrieve_trait_sample_data
 
 
 def compute_sum(rhs_val: int, lhs_val: int) -> int:
@@ -28,7 +29,8 @@ def get_trait_sample_data(trait_dataset, trait_name) -> dict:
     try to fetch from the traits dataset redis is only  used for\
     temp dataset type which is not used in this case """
 
-    sample_results = trait_dataset.retrieve_sample_data(trait_name)
+    sample_results = retrieve_trait_sample_data(
+        trait_dataset.id, trait_dataset.type, trait_name)
 
     trait_data = {}
 
diff --git a/tests/unit/computations/test_trait.py b/tests/unit/computations/test_trait.py
index 20f0546..f662325 100644
--- a/tests/unit/computations/test_trait.py
+++ b/tests/unit/computations/test_trait.py
@@ -37,7 +37,8 @@ class TestTrait(TestCase):
         self.assertEqual(results, expected_trait)
         get_sample_data.assert_called_once_with(None, "AXFDSF_AT")
 
-    def test_get_trait_sample_data(self):
+    @mock.patch("gn3.computations.traits.retrieve_trait_sample_data")
+    def test_get_trait_sample_data(self, mock_retrieve_sample_data):
         """test for getting sample data from  either\
         the trait's dataset or form redis"""
 
@@ -47,7 +48,7 @@ class TestTrait(TestCase):
             ('A/J', 7.596, None, None, 'A/J'),
             ('AKR/J', 7.774, None, None, 'AKR/J'),
             ('B6D2F1', 7.707, None, None, 'B6D2F1')]
-        trait_dataset.retrieve_sample_data.return_value = dataset_trait_sample_data
+        mock_retrieve_sample_data.return_value = dataset_trait_sample_data
 
         trait_name = "1426679_at"