fix for fetching dataset traits data

2 files changed, 12 insertions, 4 deletions

diff --git a/tests/unit/computations/test_datasets.py b/tests/unit/computations/test_datasets.py
index b696f70..44ff527 100644
--- a/tests/unit/computations/test_datasets.py
+++ b/tests/unit/computations/test_datasets.py
@@ -29,6 +29,8 @@ class TestDatasets(TestCase):
         dataset_id = "HC_M2_0606_P&"
         dataset_type = "Publish"
 
+        database = mock.Mock()
+
         dataset = {
             "id": dataset_id,
             "type": dataset_type,
@@ -40,7 +42,7 @@ class TestDatasets(TestCase):
         mock_fetch_sample_results.return_value = fetch_results
 
         results = retrieve_trait_sample_data(
-            dataset, trait_name)
+            dataset, trait_name, database)
         self.assertEqual(mock_fetch_sample_results.call_count, 1)
         self.assertEqual(results, fetch_results)
 
diff --git a/tests/unit/computations/test_trait.py b/tests/unit/computations/test_trait.py
index f662325..49c08d4 100644
--- a/tests/unit/computations/test_trait.py
+++ b/tests/unit/computations/test_trait.py
@@ -25,6 +25,8 @@ class TestTrait(TestCase):
             "WQC": 11.1
         }
 
+        database = mock.Mock()
+
         get_sample_data.return_value = expected_sample_data
 
         expected_trait = {
@@ -32,10 +34,12 @@ class TestTrait(TestCase):
             "dataset": None,
             "trait_data": expected_sample_data
         }
-        results = fetch_trait(dataset=None, trait_name="AXFDSF_AT")
+        results = fetch_trait(dataset=None,
+                              trait_name="AXFDSF_AT",
+                              database=database)
 
         self.assertEqual(results, expected_trait)
-        get_sample_data.assert_called_once_with(None, "AXFDSF_AT")
+        self.assertEqual(get_sample_data.call_count, 1)
 
     @mock.patch("gn3.computations.traits.retrieve_trait_sample_data")
     def test_get_trait_sample_data(self, mock_retrieve_sample_data):
@@ -52,8 +56,10 @@ class TestTrait(TestCase):
 
         trait_name = "1426679_at"
 
+        database = mock.Mock()
+
         results = get_trait_sample_data(
-            trait_dataset, trait_name)
+            trait_dataset, trait_name, database)
 
         expected_results = {
             "129S1/SvImJ": 7.433,