diff options
| -rw-r--r-- | gn3/db/traits.py | 18 | ||||
| -rw-r--r-- | tests/unit/db/test_traits.py | 22 |
2 files changed, 40 insertions, 0 deletions
diff --git a/gn3/db/traits.py b/gn3/db/traits.py index 4860a07..37b111e 100644 --- a/gn3/db/traits.py +++ b/gn3/db/traits.py @@ -90,3 +90,21 @@ def insert_publication(pubmed_id: int, publication: Optional[Dict], ", ".join(['%s'] * len(publication)))) with conn.cursor() as cursor: cursor.execute(insert_query, tuple(publication.values())) + +def retrieve_probeset_trait_name(threshold, name, connection): + """ + Retrieve the name for a Probeset trait + + This is extracted from the `webqtlDataset.retrieveName` function, + specifically the section dealing with 'ProbeSet' type traits + https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlDataset.py#L140-154""" + query = ( + 'SELECT Id, Name, FullName, ShortName, DataScale ' + 'FROM ProbeSetFreeze ' + 'WHERE ' + 'public > %(threshold)s ' + 'AND ' + '(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)') + with connection.cursor() as cursor: + cursor.execute(query, {"threshold": threshold, "name": name}) + return cursor.fetchone() diff --git a/tests/unit/db/test_traits.py b/tests/unit/db/test_traits.py new file mode 100644 index 0000000..6d2ba4d --- /dev/null +++ b/tests/unit/db/test_traits.py @@ -0,0 +1,22 @@ +"""Tests for gn3/db/traits.py""" +from unittest import mock, TestCase +from gn3.db.traits import retrieve_probeset_trait_name + +class TestTraitsDBFunctions(TestCase): + "Test cases for traits functions" + + def test_retrieve_probeset_trait_name(self): + """Test that the function is called correctly.""" + db_mock = mock.MagicMock() + with db_mock.cursor() as cursor: + cursor.fetchone.return_value = ( + "testName", "testNameFull", "testNameShort", "dataScale") + self.assertEqual( + retrieve_probeset_trait_name(9, "testName", db_mock), + ("testName", "testNameFull", "testNameShort", "dataScale")) + cursor.execute.assert_called_once_with( + "SELECT Id, Name, FullName, ShortName, DataScale " + "FROM ProbeSetFreeze " + "WHERE public > %(threshold)s AND " + "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)", + {"threshold": 9, "name": "testName"}) |