Merge pull request #32 from genenetwork/heatmap_decompose_db_retrieval

Heatmap decompose db retrieval

2 files changed, 269 insertions, 57 deletions

diff --git a/tests/unit/db/test_datasets.py b/tests/unit/db/test_datasets.py
new file mode 100644
index 0000000..38de0e2
--- /dev/null
+++ b/tests/unit/db/test_datasets.py
@@ -0,0 +1,133 @@
+"""Tests for gn3/db/datasets.py"""
+
+from unittest import mock, TestCase
+from gn3.db.datasets import (
+    retrieve_dataset_name,
+    retrieve_riset_fields,
+    retrieve_geno_riset_fields,
+    retrieve_publish_riset_fields,
+    retrieve_probeset_riset_fields)
+
+class TestDatasetsDBFunctions(TestCase):
+    """Test cases for datasets functions."""
+
+    def test_retrieve_dataset_name(self):
+        """Test that the function is called correctly."""
+        for trait_type, thresh, trait_name, dataset_name, columns, table in [
+                ["ProbeSet", 9, "probesetTraitName", "probesetDatasetName",
+                 "Id, Name, FullName, ShortName, DataScale", "ProbeSetFreeze"],
+                ["Geno", 3, "genoTraitName", "genoDatasetName",
+                 "Id, Name, FullName, ShortName", "GenoFreeze"],
+                ["Publish", 6, "publishTraitName", "publishDatasetName",
+                 "Id, Name, FullName, ShortName", "PublishFreeze"],
+                ["Temp", 4, "tempTraitName", "tempTraitName",
+                 "Id, Name, FullName, ShortName", "TempFreeze"]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_type=trait_type):
+                with db_mock.cursor() as cursor:
+                    cursor.fetchone.return_value = {}
+                    self.assertEqual(
+                        retrieve_dataset_name(
+                            trait_type, thresh, trait_name, dataset_name, db_mock),
+                        {})
+                    cursor.execute.assert_called_once_with(
+                        "SELECT {cols} "
+                        "FROM {table} "
+                        "WHERE public > %(threshold)s AND "
+                        "(Name = %(name)s "
+                        "OR FullName = %(name)s "
+                        "OR ShortName = %(name)s)".format(
+                            table=table, cols=columns),
+                        {"threshold": thresh, "name": dataset_name})
+
+    def test_retrieve_probeset_riset_fields(self):
+        """
+        Test that the `riset` and `riset_id` fields are retrieved appropriately
+        for the 'ProbeSet' trait type.
+        """
+        for trait_name, expected in [
+                ["testProbeSetName", {}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_name=trait_name, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ()
+                    self.assertEqual(
+                        retrieve_probeset_riset_fields(trait_name, db_mock),
+                        expected)
+                    cursor.execute.assert_called_once_with(
+                        (
+                            "SELECT InbredSet.Name, InbredSet.Id"
+                            " FROM InbredSet, ProbeSetFreeze, ProbeFreeze"
+                            " WHERE ProbeFreeze.InbredSetId = InbredSet.Id"
+                            " AND ProbeFreeze.Id = ProbeSetFreeze.ProbeFreezeId"
+                            " AND ProbeSetFreeze.Name = %(name)s"),
+                        {"name": trait_name})
+
+    def test_retrieve_riset_fields(self):
+        """
+        Test that the riset fields are set up correctly for the different trait
+        types.
+        """
+        for trait_type, trait_name, dataset_info, expected in [
+                ["Publish", "pubTraitName01", {"dataset_name": "pubDBName01"},
+                 {"dataset_name": "pubDBName01", "riset": ""}],
+                ["ProbeSet", "prbTraitName01", {"dataset_name": "prbDBName01"},
+                 {"dataset_name": "prbDBName01", "riset": ""}],
+                ["Geno", "genoTraitName01", {"dataset_name": "genoDBName01"},
+                 {"dataset_name": "genoDBName01", "riset": ""}],
+                ["Temp", "tempTraitName01", {}, {"riset": ""}],
+                ]:
+            db_mock = mock.MagicMock()
+            with self.subTest(
+                    trait_type=trait_type, trait_name=trait_name,
+                    dataset_info=dataset_info):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ("riset_name", 0)
+                    self.assertEqual(
+                        retrieve_riset_fields(
+                            trait_type, trait_name, dataset_info, db_mock),
+                        expected)
+
+    def test_retrieve_publish_riset_fields(self):
+        """
+        Test that the `riset` and `riset_id` fields are retrieved appropriately
+        for the 'Publish' trait type.
+        """
+        for trait_name, expected in [
+                ["testPublishName", {}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_name=trait_name, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ()
+                    self.assertEqual(
+                        retrieve_publish_riset_fields(trait_name, db_mock),
+                        expected)
+                    cursor.execute.assert_called_once_with(
+                        (
+                            "SELECT InbredSet.Name, InbredSet.Id"
+                            " FROM InbredSet, PublishFreeze"
+                            " WHERE PublishFreeze.InbredSetId = InbredSet.Id"
+                            " AND PublishFreeze.Name = %(name)s"),
+                        {"name": trait_name})
+
+    def test_retrieve_geno_riset_fields(self):
+        """
+        Test that the `riset` and `riset_id` fields are retrieved appropriately
+        for the 'Geno' trait type.
+        """
+        for trait_name, expected in [
+                ["testGenoName", {}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_name=trait_name, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ()
+                    self.assertEqual(
+                        retrieve_geno_riset_fields(trait_name, db_mock),
+                        expected)
+                    cursor.execute.assert_called_once_with(
+                        (
+                            "SELECT InbredSet.Name, InbredSet.Id"
+                            " FROM InbredSet, GenoFreeze"
+                            " WHERE GenoFreeze.InbredSetId = InbredSet.Id"
+                            " AND GenoFreeze.Name = %(name)s"),
+                        {"name": trait_name})
diff --git a/tests/unit/db/test_traits.py b/tests/unit/db/test_traits.py
index c8f28b5..ee98893 100644
--- a/tests/unit/db/test_traits.py
+++ b/tests/unit/db/test_traits.py
@@ -1,51 +1,20 @@
 """Tests for gn3/db/traits.py"""
 from unittest import mock, TestCase
 from gn3.db.traits import (
-    GENO_TRAIT_INFO_QUERY,
-    TEMP_TRAIT_INFO_QUERY,
-    PUBLISH_TRAIT_INFO_QUERY,
-    PROBESET_TRAIT_INFO_QUERY)
-from gn3.db.traits import (
+    build_trait_name,
+    set_haveinfo_field,
+    update_sample_data,
     retrieve_trait_info,
+    set_confidential_field,
+    set_homologene_id_field,
     retrieve_geno_trait_info,
     retrieve_temp_trait_info,
-    retrieve_trait_dataset_name,
     retrieve_publish_trait_info,
-    retrieve_probeset_trait_info,
-    update_sample_data)
-
+    retrieve_probeset_trait_info)
 
 class TestTraitsDBFunctions(TestCase):
     "Test cases for traits functions"
 
-    def test_retrieve_trait_dataset_name(self):
-        """Test that the function is called correctly."""
-        for trait_type, thresh, trait_dataset_name, columns in [
-                ["ProbeSet", 9, "testName",
-                 "Id, Name, FullName, ShortName, DataScale"],
-                ["Geno", 3, "genoTraitName", "Id, Name, FullName, ShortName"],
-                ["Publish", 6, "publishTraitName",
-                 "Id, Name, FullName, ShortName"],
-                ["Temp", 4, "tempTraitName", "Id, Name, FullName, ShortName"]]:
-            db_mock = mock.MagicMock()
-            with self.subTest(trait_type=trait_type):
-                with db_mock.cursor() as cursor:
-                    cursor.fetchone.return_value = (
-                        "testName", "testNameFull", "testNameShort",
-                        "dataScale")
-                    self.assertEqual(
-                        retrieve_trait_dataset_name(
-                            trait_type, thresh, trait_dataset_name, db_mock),
-                        ("testName", "testNameFull", "testNameShort",
-                         "dataScale"))
-                    cursor.execute.assert_called_once_with(
-                        "SELECT {cols} "
-                        "FROM {ttype}Freeze "
-                        "WHERE public > %(threshold)s AND "
-                        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)".format(
-                            cols=columns, ttype=trait_type),
-                        {"threshold": thresh, "name": trait_dataset_name})
-
     def test_retrieve_publish_trait_info(self):
         """Test retrieval of type `Publish` traits."""
         db_mock = mock.MagicMock()
@@ -54,12 +23,32 @@ class TestTraitsDBFunctions(TestCase):
             trait_source = {
                 "trait_name": "PublishTraitName", "trait_dataset_id": 1}
             self.assertEqual(
-                retrieve_publish_trait_info(
-                    trait_source,
-                    db_mock),
-                tuple())
+                retrieve_publish_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                PUBLISH_TRAIT_INFO_QUERY, trait_source)
+                ("SELECT "
+                 "PublishXRef.Id, Publication.PubMed_ID,"
+                 " Phenotype.Pre_publication_description,"
+                 " Phenotype.Post_publication_description,"
+                 " Phenotype.Original_description,"
+                 " Phenotype.Pre_publication_abbreviation,"
+                 " Phenotype.Post_publication_abbreviation,"
+                 " Phenotype.Lab_code, Phenotype.Submitter, Phenotype.Owner,"
+                 " Phenotype.Authorized_Users,"
+                 " CAST(Publication.Authors AS BINARY),"
+                 " Publication.Title, Publication.Abstract,"
+                 " Publication.Journal,"
+                 " Publication.Volume, Publication.Pages, Publication.Month,"
+                 " Publication.Year, PublishXRef.Sequence, Phenotype.Units,"
+                 " PublishXRef.comments"
+                 " FROM"
+                 " PublishXRef, Publication, Phenotype, PublishFreeze"
+                 " WHERE"
+                 " PublishXRef.Id = %(trait_name)s"
+                 " AND Phenotype.Id = PublishXRef.PhenotypeId"
+                 " AND Publication.Id = PublishXRef.PublicationId"
+                 " AND PublishXRef.InbredSetId = PublishFreeze.InbredSetId"
+                 " AND PublishFreeze.Id =%(trait_dataset_id)s"),
+                trait_source)
 
     def test_retrieve_probeset_trait_info(self):
         """Test retrieval of type `Probeset` traits."""
@@ -70,9 +59,31 @@ class TestTraitsDBFunctions(TestCase):
                 "trait_name": "ProbeSetTraitName",
                 "trait_dataset_name": "ProbeSetDatasetTraitName"}
             self.assertEqual(
-                retrieve_probeset_trait_info(trait_source, db_mock), tuple())
+                retrieve_probeset_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                PROBESET_TRAIT_INFO_QUERY, trait_source)
+                (
+                    "SELECT "
+                    "ProbeSet.name, ProbeSet.symbol, ProbeSet.description, "
+                    "ProbeSet.probe_target_description, ProbeSet.chr, "
+                    "ProbeSet.mb, ProbeSet.alias, ProbeSet.geneid, "
+                    "ProbeSet.genbankid, ProbeSet.unigeneid, ProbeSet.omim, "
+                    "ProbeSet.refseq_transcriptid, ProbeSet.blatseq, "
+                    "ProbeSet.targetseq, ProbeSet.chipid, ProbeSet.comments, "
+                    "ProbeSet.strand_probe, ProbeSet.strand_gene, "
+                    "ProbeSet.probe_set_target_region, ProbeSet.proteinid, "
+                    "ProbeSet.probe_set_specificity, "
+                    "ProbeSet.probe_set_blat_score, "
+                    "ProbeSet.probe_set_blat_mb_start, "
+                    "ProbeSet.probe_set_blat_mb_end, "
+                    "ProbeSet.probe_set_strand, ProbeSet.probe_set_note_by_rw, "
+                    "ProbeSet.flag "
+                    "FROM "
+                    "ProbeSet, ProbeSetFreeze, ProbeSetXRef "
+                    "WHERE "
+                    "ProbeSetXRef.ProbeSetFreezeId = ProbeSetFreeze.Id "
+                    "AND ProbeSetXRef.ProbeSetId = ProbeSet.Id "
+                    "AND ProbeSetFreeze.Name = %(trait_dataset_name)s "
+                    "AND ProbeSet.Name = %(trait_name)s"), trait_source)
 
     def test_retrieve_geno_trait_info(self):
         """Test retrieval of type `Geno` traits."""
@@ -83,9 +94,19 @@ class TestTraitsDBFunctions(TestCase):
                 "trait_name": "GenoTraitName",
                 "trait_dataset_name": "GenoDatasetTraitName"}
             self.assertEqual(
-                retrieve_geno_trait_info(trait_source, db_mock), tuple())
+                retrieve_geno_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                GENO_TRAIT_INFO_QUERY, trait_source)
+                (
+                    "SELECT "
+                    "Geno.name, Geno.chr, Geno.mb, Geno.source2, Geno.sequence "
+                    "FROM "
+                    "Geno, GenoFreeze, GenoXRef "
+                    "WHERE "
+                    "GenoXRef.GenoFreezeId = GenoFreeze.Id "
+                    "AND GenoXRef.GenoId = Geno.Id "
+                    "AND GenoFreeze.Name = %(trait_dataset_name)s "
+                    "AND Geno.Name = %(trait_name)s"),
+                trait_source)
 
     def test_retrieve_temp_trait_info(self):
         """Test retrieval of type `Temp` traits."""
@@ -94,26 +115,51 @@ class TestTraitsDBFunctions(TestCase):
             cursor.fetchone.return_value = tuple()
             trait_source = {"trait_name": "TempTraitName"}
             self.assertEqual(
-                retrieve_temp_trait_info(trait_source, db_mock), tuple())
+                retrieve_temp_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                TEMP_TRAIT_INFO_QUERY, trait_source)
+                "SELECT name, description FROM Temp WHERE Name = %(trait_name)s",
+                trait_source)
+
+    def test_build_trait_name_with_good_fullnames(self):
+        """
+        Check that the name is built correctly.
+        """
+        for fullname, expected in [
+                ["testdb::testname",
+                 {"db": {"dataset_name": "testdb", "dataset_type": "ProbeSet"},
+                  "trait_name": "testname", "cellid": "",
+                  "trait_fullname": "testdb::testname"}],
+                ["testdb::testname::testcell",
+                 {"db": {"dataset_name": "testdb", "dataset_type": "ProbeSet"},
+                  "trait_name": "testname", "cellid": "testcell",
+                  "trait_fullname": "testdb::testname::testcell"}]]:
+            with self.subTest(fullname=fullname):
+                self.assertEqual(build_trait_name(fullname), expected)
+
+    def test_build_trait_name_with_bad_fullnames(self):
+        """
+        Check that an exception is raised if the full name format is wrong.
+        """
+        for fullname in ["", "test", "test:test"]:
+            with self.subTest(fullname=fullname):
+                with self.assertRaises(AssertionError, msg="Name format error"):
+                    build_trait_name(fullname)
 
     def test_retrieve_trait_info(self):
         """Test that information on traits is retrieved as appropriate."""
-        for trait_type, trait_name, trait_dataset_id, trait_dataset_name, in [
-                ["Publish", "PublishTraitName", 1, "PublishDatasetTraitName"],
-                ["ProbeSet", "ProbeSetTraitName", 2, "ProbeSetDatasetTraitName"],
-                ["Geno", "GenoTraitName", 3, "GenoDatasetTraitName"],
-                ["Temp", "TempTraitName", 4, "TempDatasetTraitName"]]:
+        for threshold, trait_fullname, expected in [
+                [9, "pubDb::PublishTraitName::pubCell", {"haveinfo": 0}],
+                [5, "prbDb::ProbeSetTraitName::prbCell", {"haveinfo": 0}],
+                [12, "genDb::GenoTraitName", {"haveinfo": 0}],
+                [6, "tmpDb::TempTraitName", {"haveinfo": 0}]]:
             db_mock = mock.MagicMock()
-            with self.subTest(trait_type=trait_type):
+            with self.subTest(trait_fullname=trait_fullname):
                 with db_mock.cursor() as cursor:
                     cursor.fetchone.return_value = tuple()
                     self.assertEqual(
                         retrieve_trait_info(
-                            trait_type, trait_name, trait_dataset_id,
-                            trait_dataset_name, db_mock),
-                        tuple())
+                            threshold, trait_fullname, db_mock),
+                        expected)
 
     def test_update_sample_data(self):
         """Test that the SQL queries when calling update_sample_data are called with
@@ -143,3 +189,36 @@ class TestTraitsDBFunctions(TestCase):
                  mock.call(PUBLISH_SE_SQL, (2.3, 10, 8967049)),
                  mock.call(N_STRAIN_SQL, (2, 10, 8967049))]
             )
+
+    def test_set_haveinfo_field(self):
+        """Test that the `haveinfo` field is set up correctly"""
+        for trait_info, expected in [
+                [{}, {"haveinfo": 0}],
+                [{"k1": "v1"}, {"k1": "v1", "haveinfo": 1}]]:
+            with self.subTest(trait_info=trait_info, expected=expected):
+                self.assertEqual(set_haveinfo_field(trait_info), expected)
+
+    def test_set_homologene_id_field(self):
+        """Test that the `homologene_id` field is set up correctly"""
+        for trait_type, trait_info, expected in [
+                ["Publish", {}, {"homologeneid": None}],
+                ["ProbeSet", {}, {"homologeneid": None}],
+                ["Geno", {}, {"homologeneid": None}],
+                ["Temp", {}, {"homologeneid": None}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_info=trait_info, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.fetchone.return_value = ()
+                    self.assertEqual(
+                        set_homologene_id_field(trait_type, trait_info, db_mock), expected)
+
+    def test_set_confidential_field(self):
+        """Test that the `confidential` field is set up correctly"""
+        for trait_type, trait_info, expected in [
+                ["Publish", {}, {"confidential": 0}],
+                ["ProbeSet", {}, {}],
+                ["Geno", {}, {}],
+                ["Temp", {}, {}]]:
+            with self.subTest(trait_info=trait_info, expected=expected):
+                self.assertEqual(
+                    set_confidential_field(trait_type, trait_info), expected)