From fc233942e9118a341001f3357a5d2b1c65187736 Mon Sep 17 00:00:00 2001
From: Muriithi Frederick Muriuki
Date: Tue, 17 Aug 2021 08:47:11 +0300
Subject: Add tests and fix errors caught with tests

Issue:
https://github.com/genenetwork/gn-gemtext-threads/blob/main/topics/gn1-migration-to-gn2/clustering.gmi

* gn3/computations/heatmap.py: fix errors
* tests/unit/computations/test_heatmap.py: new tests

  Add new tests with the expected source data format, and expected results.
  Fix all errors that were caught by running the tests
---
 gn3/computations/heatmap.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

(limited to 'gn3')

diff --git a/gn3/computations/heatmap.py b/gn3/computations/heatmap.py
index a0e778a..8a86fe8 100644
--- a/gn3/computations/heatmap.py
+++ b/gn3/computations/heatmap.py
@@ -34,11 +34,11 @@ def export_trait_data(
     """
     def __export_all_types(tdata, strain):
         sample_data = []
-        if tdata[strain]["val"]:
-            sample_data.append(tdata[strain]["val"])
+        if tdata[strain]["value"]:
+            sample_data.append(tdata[strain]["value"])
             if var_exists:
-                if tdata[strain].var:
-                    sample_data.append(tdata[strain]["var"])
+                if tdata[strain]["variance"]:
+                    sample_data.append(tdata[strain]["variance"])
                 else:
                     sample_data.append(None)
             if n_exists:
@@ -58,15 +58,15 @@ def export_trait_data(
 
     def __exporter(accumulator, strain):
         # pylint: disable=[R0911]
-        if trait_data.has_key(strain):
+        if strain in trait_data["data"]:
             if dtype == "val":
-                return accumulator + (trait_data[strain]["val"], )
+                return accumulator + (trait_data["data"][strain]["value"], )
             if dtype == "var":
-                return accumulator + (trait_data[strain]["var"], )
+                return accumulator + (trait_data["data"][strain]["variance"], )
             if dtype == "N":
-                return trait_data[strain]["ndata"]
+                return accumulator + (trait_data["data"][strain]["ndata"], )
             if dtype == "all":
-                return accumulator + __export_all_types(trait_data, strain)
+                return accumulator + __export_all_types(trait_data["data"], strain)
             raise KeyError("Type `%s` is incorrect" % dtype)
         if var_exists and n_exists:
             return accumulator + (None, None, None)
-- 
cgit v1.2.3