add composite key + update tests

Author: korikuzma

src/anyvlm/storage/orm.py

@@ -42,13 +42,13 @@ def __tablename__(self) -> str:
         return "allele_frequency_data"
 
     vrs_id: Mapped[str] = mapped_column(String, primary_key=True)
+    consequence: Mapped[str] = mapped_column(String, primary_key=True)
     af: Mapped[float] = mapped_column(Float)
     ac: Mapped[int] = mapped_column(Integer)
     an: Mapped[int] = mapped_column(Integer)
     ac_het: Mapped[int] = mapped_column(Integer)
     ac_hom: Mapped[int] = mapped_column(Integer)
     ac_hemi: Mapped[int] = mapped_column(Integer)
-    consequence: Mapped[str] = mapped_column(String)
     filter: Mapped[list[str]] = mapped_column(ARRAY(String))
 
 

tests/unit/storage/test_postgres_unit.py

@@ -6,6 +6,7 @@
 from ga4gh.vrs.models import Allele, LiteralSequenceExpression
 from tests.conftest import return_cafs
 
+from anyvlm.storage import orm
 from anyvlm.storage.postgres import PostgresObjectStore
 
 
@@ -29,27 +30,43 @@ def test_add_allele_frequency(
     request, caf_fixture_name: str, postgres_storage: PostgresObjectStore
 ):
     """Test that add_allele_frequency method works correctly"""
-    caf = request.getfixturevalue(caf_fixture_name)
-
-    def insert_and_verify_one_row():
-        postgres_storage.add_allele_frequency(caf)
-
-        rows = return_cafs(postgres_storage)
-        assert len(rows) == 1
-
-        db_record = rows[0]
 
+    def verify_row(db_record: orm.AlleleFrequencyData, expected_consequence: str):
+        """Verify row"""
         assert db_record.vrs_id == "ga4gh:VA.J3Hi64dkKFKdnKIwB2419Qz3STDB2sJq"
         assert db_record.af == 0.000162232
         assert db_record.ac == 1
         assert db_record.an == 6164
         assert db_record.ac_het == 1
         assert db_record.ac_hom == 0
-        assert db_record.consequence == "intron_variant"
+        assert db_record.consequence == expected_consequence
         assert db_record.filter == ["LowQual", "NO_HQ_GENOTYPES"]
 
+    def insert_and_verify_rows(
+        caf_obj: CohortAlleleFrequencyStudyResult, expected_consequences: list[str]
+    ):
+        """Insert caf record and verify rows"""
+        postgres_storage.add_allele_frequency(caf_obj)
+
+        rows = return_cafs(postgres_storage)
+        assert len(rows) == len(expected_consequences)
+
+        for index, consequence in enumerate(expected_consequences):
+            db_record = rows[index]
+            verify_row(db_record, consequence)
+
+    caf = request.getfixturevalue(caf_fixture_name)
+    expected_consequences = ["intron_variant"]
+
     # first insert
-    insert_and_verify_one_row()
+    insert_and_verify_rows(caf, expected_consequences=expected_consequences)
 
     # second insert
-    insert_and_verify_one_row()
+    insert_and_verify_rows(caf, expected_consequences=expected_consequences)
+
+    # test composite key
+    new_caf = caf.model_copy(deep=True)
+    new_consequence = "missense_variant"
+    new_caf.ancillaryResults["consequence"] = new_consequence
+    expected_consequences.append(new_consequence)
+    insert_and_verify_rows(new_caf, expected_consequences=expected_consequences)