from src.data_reader import load_train_test_df
from src.feature_extractors.bert_pretrain_extractor import \
    BertPretrainFeatureExtractor


def test_pretrain_feature_extractor():
    models = ['distilbert-base-uncased-finetuned-sst-2-english', 'bert-base-uncased']
    train_df, _ = load_train_test_df(is_testing=True)

    for model_name in models:
        feature_extractor = BertPretrainFeatureExtractor(model_name=model_name)

        output_features = feature_extractor.generate_features(train_df.full_text)

        assert len(output_features) == 5 and len(output_features.columns) == 768