test_main.py - master - ilmerkul/CountVectorizer

test_main.py
46 строк2 KB
Ilya Merkulov
class was created
15 окт 2024, 09:39
15 окт 2024, 09:39d86f86b
import unittest
from main import CountVectorizer


class TestCountVectorizer(unittest.TestCase):
    def test_use_fit_transform(self):
        corpus = ['Crock Pot Pasta Never boil pasta again',
                  'Pasta Pomodoro Fresh ingredients Parmesan to taste']

        vectorizer = CountVectorizer()
        self.assertEqual(vectorizer.fit_transform(corpus), [[1, 1, 2, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                                                            [0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1]],
                         'error test_use_fit_transform.fit_transform')
        self.assertEqual(vectorizer.get_feature_names(),
                         ['crock', 'pot', 'pasta', 'never', 'boil', 'again', 'pomodoro', 'fresh', 'ingredients',
                          'parmesan', 'to', 'taste'], 'error test_use_fit_transform.get_feature_names')

    def test_use_lower_fit_transform(self):
        corpus = ['Crock Pot Pasta Never boil pasta again',
                  'Pasta Pomodoro Fresh ingredients Parmesan to taste']

        vectorizer = CountVectorizer()
        self.assertEqual(vectorizer.fit_transform(corpus, lowercase=False),
                         [[1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]],
                         'error count_matrix')
        self.assertEqual(vectorizer.get_feature_names(),
                         ['Crock', 'Pot', 'Pasta', 'Never', 'boil', 'pasta', 'again', 'Pomodoro', 'Fresh',
                          'ingredients',
                          'Parmesan', 'to', 'taste'], 'error test_use_lower_fit_transform.get_feature_names')

    def test_use_fit(self):
        corpus = ['Crock Pot Pasta Never boil pasta again',
                  'Pasta Pomodoro Fresh ingredients Parmesan to taste']

        vectorizer = CountVectorizer()
        vectorizer.fit(corpus)
        self.assertEqual(vectorizer.get_feature_names(),
                         ['crock', 'pot', 'pasta', 'never', 'boil', 'again', 'pomodoro', 'fresh', 'ingredients',
                          'parmesan', 'to', 'taste'], 'error test_use_fit.get_feature_names')
        self.assertEqual(vectorizer.transform(['boil taste', 'pasta Fresh']), [[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1],
                                                                         [0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0]],
                         'error test_use_fit.fit')


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