import os
import shutil
import tempfile
import unittest
from collections import Counter
from pathlib import Path
import pkg_resources
from medacy.data.dataset import Dataset
from medacy.data.annotations import Annotations
from medacy.data.data_file import DataFile
from medacy.tests.sample_data import test_dir
class TestDataset(unittest.TestCase):
"""Unit tests for Dataset"""
@classmethod
def setUpClass(cls):
cls.dataset = Dataset(os.path.join(test_dir, 'sample_dataset_1'))
cls.prediction_directory = tempfile.mkdtemp() # Set up predict directory
cls.entities = cls.dataset.get_labels(as_list=True)
cls.ann_files = []
# Fill directory of prediction files (only the text files)
for data_file in cls.dataset:
new_file_path = os.path.join(cls.prediction_directory, data_file.file_name + '.txt')
shutil.copyfile(data_file.txt_path, new_file_path)
# Fill a directory with just ann files
cls.ann_dir = tempfile.mkdtemp()
for data_file in cls.dataset:
new_ann_path = os.path.join(cls.ann_dir, data_file.file_name + '.ann')
shutil.copyfile(data_file.ann_path, new_ann_path)
@classmethod
def tearDownClass(cls):
pkg_resources.cleanup_resources()
for directory in [cls.prediction_directory, cls.ann_dir]:
shutil.rmtree(directory)
def test_init(self):
"""Tests initializing Datasets from different directories to see that they create accurate DataFiles"""
# Test both txt, ann, and metamapped
test_dir_path = Path(self.dataset.data_directory)
expected = [
DataFile(
file_name="PMC1257590",
txt_path=test_dir_path / "PMC1257590.txt",
ann_path=test_dir_path / "PMC1257590.ann",
metamapped_path=test_dir_path / "metamapped" / "PMC1257590.metamapped"
),
DataFile(
file_name="PMC1314908",
txt_path=test_dir_path / "PMC1314908.txt",
ann_path=test_dir_path / "PMC1314908.ann",
metamapped_path=test_dir_path / "metamapped" / "PMC1314908.metamapped"
),
DataFile(
file_name="PMC1392236",
txt_path=test_dir_path / "PMC1392236.txt",
ann_path=test_dir_path / "PMC1392236.ann",
metamapped_path=test_dir_path / "metamapped" / "PMC1392236.metamapped"
)
]
expected.sort(key=lambda x: x.file_name)
actual = list(self.dataset)
self.assertListEqual(actual, expected)
# Test txt only
test_dir_path = Path(self.prediction_directory)
expected = [
DataFile(
file_name="PMC1257590",
txt_path=test_dir_path / "PMC1257590.txt",
ann_path=None,
metamapped_path=None
),
DataFile(
file_name="PMC1314908",
txt_path=test_dir_path / "PMC1314908.txt",
ann_path=None,
metamapped_path=None
),
DataFile(
file_name="PMC1392236",
txt_path=test_dir_path / "PMC1392236.txt",
ann_path=None,
metamapped_path=None
)
]
expected.sort(key=lambda x: x.file_name)
actual = list(Dataset(self.prediction_directory))
self.assertListEqual(actual, expected)
# Test ann only
test_dir_path = Path(self.ann_dir)
expected = [
DataFile(
file_name="PMC1257590",
txt_path=None,
ann_path=test_dir_path / "PMC1257590.ann",
metamapped_path=None
),
DataFile(
file_name="PMC1314908",
txt_path=None,
ann_path=test_dir_path / "PMC1314908.ann",
metamapped_path=None,
),
DataFile(
file_name="PMC1392236",
txt_path=None,
ann_path=test_dir_path / "PMC1392236.ann",
metamapped_path=None
)
]
expected.sort(key=lambda x: x.file_name)
actual = list(Dataset(self.ann_dir))
self.assertListEqual(actual, expected)
def test_init_with_data_limit(self):
"""Tests that initializing with a data limit works"""
dataset = Dataset(self.dataset.data_directory, data_limit=1)
self.assertEqual(len(list(dataset)), 1)
def test_generate_annotations(self):
"""Tests that generate_annotations() creates Annotations objects"""
for ann in self.dataset.generate_annotations():
self.assertIsInstance(ann, Annotations)
def test_get_labels(self):
"""Tests that get_labels returns a set of the correct labels"""
expected = {
'DoseFrequency', 'SampleSize', 'TimeUnits', 'Vehicle', 'TestArticlePurity', 'Endpoint', 'TestArticle',
'GroupName', 'DoseDurationUnits', 'GroupSize', 'TimeAtFirstDose', 'Dose', 'DoseDuration', 'Species',
'DoseUnits', 'Sex', 'EndpointUnitOfMeasure', 'TimeEndpointAssessed', 'DoseRoute', 'CellLine', 'Strain'
}
actual = self.dataset.get_labels()
self.assertSetEqual(actual, expected)
def test_compute_counts(self):
"""Tests that compute_counts() returns a Counter containing counts for all labels"""
counts = self.dataset.compute_counts()
self.assertIsInstance(counts, Counter)
for label in self.dataset.get_labels():
self.assertIn(label, counts.keys())
def test_getitem(self):
"""Tests that some_dataset['filename'] returns an Annotations for 'filename.ann', or raises FileNotFoundError"""
some_file_name = self.dataset.data_files[0].file_name
result = self.dataset[some_file_name]
self.assertIsInstance(result, Annotations)
with self.assertRaises(FileNotFoundError):
ann = self.dataset['notafilepath']
def test_valid_datafiles(self):
"""Tests that each DataFile in the Dataset is an existing file"""
for d in self.dataset:
self.assertTrue(os.path.isfile(d.txt_path))
self.assertTrue(os.path.isfile(d.ann_path))
self.assertTrue(os.path.isfile(d.metamapped_path))
if __name__ == '__main__':
unittest.main()
