def test_transform(self):
"""Test NVDFeedPreprocessor `transform` method."""
# custom parameters
prep = NVDFeedPreprocessor(
attributes=TEST_CVE_ATTR # only extract cve_id
)
result, = prep.transform(X=[TEST_CVE])
# result should be a tuple of default handlers and cve attributes
# check only the cve attributes here to avoid calling handler
# separately
self.assertEqual(result[-len(TEST_CVE_ATTR):], TEST_CVE_ATTR_VALS)
def test_fit_transform(self):
"""Test LabelPreprocessor `fit_transform` method."""
hook = Hook(key='label', func=lambda p, d: 'label')
attributes = ['project', 'description']
test_data = [TEST_CVE]
# prepare the data for label preprocessor
feed_prep = NVDFeedPreprocessor(attributes, skip_duplicity=True)
test_data = feed_prep.transform(X=test_data)
label_prep = LabelPreprocessor(feed_attributes=attributes,
output_attributes=['description'],
hook=hook)
result = label_prep.fit_transform(test_data)
# check not empty
self.assertTrue(any(result))
# check that correct label is returned by the Hook
label = result[0].label
self.assertEqual(label, 'label')