Update tests

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

* Fix tests to take current changes into consideration.

2 files changed, 57 insertions, 72 deletions

diff --git a/tests/unit/db/test_datasets.py b/tests/unit/db/test_datasets.py
index 34fe7f0..4f405cb 100644
--- a/tests/unit/db/test_datasets.py
+++ b/tests/unit/db/test_datasets.py
@@ -1,51 +1,56 @@
 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):
 
-    def test_retrieve_trait_dataset_name(self):
+    def test_retrieve_dataset_name(self):
         """Test that the function is called correctly."""
-        for trait_type, thresh, trait_dataset_name, columns, table in [
-                ["ProbeSet", 9, "testName",
+        for trait_type, thresh, trait_name, dataset_name, columns, table in [
+                ["ProbeSet", 9, "probesetTraitName", "probesetDatasetName",
                  "Id, Name, FullName, ShortName, DataScale", "ProbeSetFreeze"],
-                ["Geno", 3, "genoTraitName", "Id, Name, FullName, ShortName",
-                 "GenoFreeze"],
-                ["Publish", 6, "publishTraitName",
+                ["Geno", 3, "genoTraitName", "genoDatasetName",
+                 "Id, Name, FullName, ShortName", "GenoFreeze"],
+                ["Publish", 6, "publishTraitName", "publishDatasetName",
                  "Id, Name, FullName, ShortName", "PublishFreeze"],
-                ["Temp", 4, "tempTraitName", "Id, Name, FullName, ShortName",
-                 "TempFreeze"]]:
+                ["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 = (
-                        "testName", "testNameFull", "testNameShort",
-                        "dataScale")
+                    cursor.fetchone.return_value = {}
                     self.assertEqual(
-                        retrieve_trait_dataset_name(
-                            trait_type, thresh, trait_dataset_name, db_mock),
-                        ("testName", "testNameFull", "testNameShort",
-                         "dataScale"))
+                        retrieve_dataset_name(
+                            trait_type, thresh, trait_name, dataset_name, db_mock),
+                        {})
                     cursor.execute.assert_called_once_with(
-                        "SELECT %(columns)s "
-                        "FROM %(table)s "
+                        "SELECT {cols} "
+                        "FROM {table} "
                         "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,
-                         "table": table, "columns": columns})
+                        "(Name = %(name)s "
+                        "OR FullName = %(name)s "
+                        "OR ShortName = %(name)s)".format(
+                            table=table, cols=columns, ttype=trait_type),
+                        {"threshold": thresh, "name": dataset_name})
 
-    def test_set_probeset_riset_fields(self):
+    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", ()]]:
+                ["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(
-                        set_probeset_riset_fields(trait_name, db_mock), expected)
+                        retrieve_probeset_riset_fields(trait_name, db_mock),
+                        expected)
                     cursor.execute.assert_called_once_with(
                         (
                             "SELECT InbredSet.Name, InbredSet.Id"
@@ -55,54 +60,45 @@ class TestDatasetsDBFunctions(TestCase):
                             " AND ProbeSetFreeze.Name = %(name)s"),
                         {"name": trait_name})
 
-    def test_set_riset_fields(self):
+    def test_retrieve_riset_fields(self):
         """
         Test that the riset fields are set up correctly for the different trait
         types.
         """
-        for trait_info, expected in [
-                [{}, {}],
-                [{"haveinfo": 0, "type": "Publish"},
-                 {"haveinfo": 0, "type": "Publish"}],
-                [{"haveinfo": 0, "type": "ProbeSet"},
-                 {"haveinfo": 0, "type": "ProbeSet"}],
-                [{"haveinfo": 0, "type": "Geno"},
-                 {"haveinfo": 0, "type": "Geno"}],
-                [{"haveinfo": 0, "type": "Temp"},
-                 {"haveinfo": 0, "type": "Temp"}],
-                [{"haveinfo": 1, "type": "Publish", "name": "test"},
-                 {"haveinfo": 1, "type": "Publish", "name": "test",
-                  "riset": "riset_name", "risetid": 0}],
-                [{"haveinfo": 1, "type": "ProbeSet", "name": "test"},
-                 {"haveinfo": 1, "type": "ProbeSet", "name": "test",
-                  "riset": "riset_name", "risetid": 0}],
-                [{"haveinfo": 1, "type": "Geno", "name": "test"},
-                 {"haveinfo": 1, "type": "Geno", "name": "test",
-                  "riset": "riset_name", "risetid": 0}],
-                [{"haveinfo": 1, "type": "Temp", "name": "test"},
-                 {"haveinfo": 1, "type": "Temp", "name": "test", "riset": None,
-                  "risetid": None}]
-        ]:
+        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_info=trait_info, expected=expected):
+            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(
-                        set_riset_fields(trait_info, db_mock), expected)
+                        retrieve_riset_fields(
+                            trait_type, trait_name, dataset_info, db_mock),
+                        expected)
 
-    def test_set_publish_riset_fields(self):
+    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", ()]]:
+                ["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(
-                        set_publish_riset_fields(trait_name, db_mock), expected)
+                        retrieve_publish_riset_fields(trait_name, db_mock),
+                        expected)
                     cursor.execute.assert_called_once_with(
                         (
                             "SELECT InbredSet.Name, InbredSet.Id"
@@ -111,19 +107,20 @@ class TestDatasetsDBFunctions(TestCase):
                             " AND PublishFreeze.Name = %(name)s"),
                         {"name": trait_name})
 
-    def test_set_geno_riset_fields(self):
+    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", ()]]:
+                ["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(
-                        set_geno_riset_fields(trait_name, db_mock), expected)
+                        retrieve_geno_riset_fields(trait_name, db_mock),
+                        expected)
                     cursor.execute.assert_called_once_with(
                         (
                             "SELECT InbredSet.Name, InbredSet.Id"
diff --git a/tests/unit/db/test_traits.py b/tests/unit/db/test_traits.py
index 7d161bf..5f52c18 100644
--- a/tests/unit/db/test_traits.py
+++ b/tests/unit/db/test_traits.py
@@ -142,25 +142,13 @@ class TestTraitsDBFunctions(TestCase):
         """Test that information on traits is retrieved as appropriate."""
         for trait_type, threshold, trait_fullname, expected in [
                 ["Publish", 9, "pubDb::PublishTraitName::pubCell",
-                 {"haveinfo": 0, "homologeneid": None, "type": "Publish",
-                  "confidential": 0, "db": {"dataset_name": "pubDb"},
-                  "trait_name": "PublishTraitName", "cellid": "pubCell",
-                  "trait_fullname": "pubDb::PublishTraitName::pubCell"}],
+                 {"haveinfo": 0}],
                 ["ProbeSet", 5, "prbDb::ProbeSetTraitName::prbCell",
-                 {"haveinfo": 0, "homologeneid": None, "type": "ProbeSet",
-                  "trait_fullname": "prbDb::ProbeSetTraitName::prbCell",
-                  "db": {"dataset_name": "prbDb"},
-                  "trait_name": "ProbeSetTraitName", "cellid": "prbCell"}],
+                 {"haveinfo": 0}],
                 ["Geno", 12, "genDb::GenoTraitName",
-                 {"haveinfo": 0, "homologeneid": None, "type": "Geno",
-                  "trait_fullname": "genDb::GenoTraitName",
-                  "db": {"dataset_name": "genDb"},
-                  "trait_name": "GenoTraitName", "cellid": ""}],
+                 {"haveinfo": 0}],
                 ["Temp", 6, "tmpDb::TempTraitName",
-                 {"haveinfo": 0, "homologeneid": None, "type": "Temp",
-                  "trait_fullname": "tmpDb::TempTraitName",
-                  "db": {"dataset_name": "tmpDb"},
-                  "trait_name": "TempTraitName", "cellid": ""}]]:
+                 {"haveinfo": 0}]]:
             db_mock = mock.MagicMock()
             with self.subTest(trait_type=trait_type):
                 with db_mock.cursor() as cursor: