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"})