From f7fb8bf6230ca9deaed382a610c40e03d74fca14 Mon Sep 17 00:00:00 2001 From: Hironsan Date: Sun, 24 Apr 2022 17:57:24 +0900 Subject: [PATCH] Add a test case for dataset --- backend/data_export/tests/test_dataset.py | 25 +++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 backend/data_export/tests/test_dataset.py diff --git a/backend/data_export/tests/test_dataset.py b/backend/data_export/tests/test_dataset.py new file mode 100644 index 00000000..f7029ebf --- /dev/null +++ b/backend/data_export/tests/test_dataset.py @@ -0,0 +1,25 @@ +import unittest +from unittest.mock import MagicMock + +import pandas as pd +from pandas.testing import assert_frame_equal + +from data_export.pipeline.dataset import Dataset + + +class TestDataset(unittest.TestCase): + def setUp(self): + example = MagicMock() + example.to_dict.return_value = {"data": "example"} + self.examples = MagicMock() + self.examples.__iter__.return_value = [example] + label = MagicMock() + label.find_by.return_value = {"labels": ["label"]} + self.labels = MagicMock() + self.labels.__iter__.return_value = [label] + + def test_to_dataframe(self): + dataset = Dataset(self.examples, self.labels) + df = dataset.to_dataframe() + expected = pd.DataFrame([{"data": "example", "labels": ["label"]}]) + assert_frame_equal(df, expected)