File size: 566 Bytes
8b414b0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import pandas as pd

from src.data_reader import load_train_test_df
from src.feature_extractors.term_frequency_feature_extractor import \
    TermFrequencyFeatureExtractor


def test_text_feature_generation():
    N_BINS = 20
    train_df, __ = load_train_test_df(is_testing=True)

    feature_extractor = TermFrequencyFeatureExtractor(n_bins=N_BINS)
    feature_df: pd.DataFrame = feature_extractor.generate_features(train_df.full_text)

    assert len(feature_df.columns) == N_BINS - 1
    assert len(feature_df) == 5
    assert not feature_df.isna().values.any()