def test_compute_metrics_adds_composite_score_column(
data_selector: DataSelector, corpus: Corpus):
expected_features = (
len(data_selector.similarity_metrics) +
len(data_selector.diversity_metrics) + 1 # composite score
)
data_selector.fit(corpus)
scores = data_selector.compute_metrics(corpus)
assert scores.shape[1] == expected_features
assert "composite" in scores
def test_compute_diversity_return_dataframe_of_correct_shape(
data_selector: DataSelector, corpus: Corpus):
data_selector.fit(corpus)
scores = data_selector.compute_diversities(corpus)
assert scores.shape[1] == len(data_selector.diversity_metrics)