import os
import sys
import shutil
import unittest

sys.path.append(os.path.join(os.path.dirname(__file__), "../"))
import chazutsu.datasets
from tests.dataset_base_test import DatasetTestCase


class TestReutersNews(DatasetTestCase):

    def test_topics(self):
        r = chazutsu.datasets.ReutersNews(kind="topics").download(directory=self.test_dir, force=True)
        self.assertTrue(len(r.data().columns), 3)
        print(r.data().columns)
        shutil.rmtree(r.root)

    def test_industries(self):
        r = chazutsu.datasets.ReutersNews().industries().download(directory=self.test_dir, force=True)
        self.assertTrue(len(r.data().columns), 3)
        print(r.data().columns)
        shutil.rmtree(r.root)

    def test_regions(self):
        r = chazutsu.datasets.ReutersNews().regions().download(directory=self.test_dir, force=True)
        self.assertTrue(len(r.data().columns), 2)
        print(r.data().columns)
        shutil.rmtree(r.root)


if __name__ == "__main__":
    unittest.main()
