Initial commit

tests/test_pc_api.py

@@ -0,0 +1,48 @@
+import polars as pl
+
+from pcdigitizer import Annotation, PubChemAPI
+
+
+def test_source_df(mock_session):
+    df = PubChemAPI.get_sources(session=mock_session)
+
+    columns = df.columns
+    assert len(columns) == 23
+    assert columns[0] == "Source Name"
+    assert columns[-3] == "License URL"
+
+    source_name = "Hazardous Substances Data Bank"
+    source_count = df.select(
+        (pl.col("Source Name").str.contains(source_name)).sum()
+    ).item()
+    assert source_count == 1
+
+
+def test_source_annotations(mock_session):
+    annotations = PubChemAPI.get_source_annotations(
+        "Hazardous Substances Data Bank (HSDB)", session=mock_session
+    )
+    assert len(annotations["Compound"]) == 246
+
+
+def test_annotations(mock_session):
+    annotations = PubChemAPI.get_annotations(session=mock_session)
+    annotation_types = sorted(annotations.keys())
+    assert annotation_types == [
+        "Assay",
+        "Cell",
+        "Compound",
+        "Element",
+        "Gene",
+        "Pathway",
+        "Protein",
+        "Taxonomy",
+    ]
+    assert "Dissociation Constants" in annotations["Compound"]
+
+
+def test_get_data(mock_session):
+    data = PubChemAPI.get_data(
+        Annotation.DISSOCIATION_CONSTANTS, 1, session=mock_session
+    )
+    assert len(data) == 1000