from pathlib import Path

import pandas as pd

import ehrapy as ep

from tests.conftest import TEST_DATA_PATH

CURRENT_DIR = Path(__file__).parent

_TEST_PATH = f"{CURRENT_DIR}/test_data_nlp"

class TestMedCAT:

    def test_add_medcat_annotation_to_obs(self):

        # created manually a small dataset with annotations to use here

        adata = ep.io.read_csv(f"{TEST_DATA_PATH}/dataset1.csv")

        adata.uns["medcat_annotations"] = pd.read_csv(f"{_TEST_PATH}/medcat_annotations1.csv")

        ep.tl.add_medcat_annotation_to_obs(adata, name="Diabetes")

        assert "Diabetes" in adata.obs.columns