def test_filter_by_handler(self):
"""Test NVDFeedPreprocessor `_filter_by_handler` method."""
prep = NVDFeedPreprocessor()
# make another cve
cve_without_ref = TestCVE()
# add reasonable reference
cve_without_ref.references = 'non-matching-reference'
cves = [cve_without_ref]
cves = prep._filter_by_handler(cves) # pylint: disable=protected-access
# check that cves list is empty
self.assertFalse(any(cves))
# noinspection PyTypeChecker
cves = [TEST_CVE]
cves = prep._filter_by_handler(cves) # pylint: disable=protected-access
self.assertTrue(any(cves))
self.assertIsInstance(cves[0], tuple)
# resulting tuple
cve, ref = cves[0]
# check that cves list is not empty
self.assertIsInstance(cve, TestCVE)
self.assertIsInstance(ref, str)
def test_get_cve_attributes(self):
"""Test NVDFeedPreprocessor `_get_cve_attributes` method."""
prep = NVDFeedPreprocessor()
cve_tuple, = prep._filter_by_handler(cves=[TestCVE()])
result = prep._get_cve_attributes(cve_tuple) # pylint: disable=protected-access
self.assertIsInstance(result, tuple)
self.assertEqual(result.repository, TEST_REPOSITORY)
self.assertIsInstance(result.references, list)
def test_init(self):
"""Test NVDFeedPreprocessor `__init__` method."""
# default parameters
prep = NVDFeedPreprocessor()
self.assertIsInstance(prep, NVDFeedPreprocessor)
# custom parameters
prep = NVDFeedPreprocessor(attributes=['cve_id'])
self.assertIsInstance(prep, NVDFeedPreprocessor)
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 NVDFeedPreprocessor `fit_transform` method."""
# default parameters
prep = NVDFeedPreprocessor()
result, = prep.fit_transform(X=[TEST_CVE])
# result should return default handler attributes
self.assertEqual(result.user, 'python')
self.assertEqual(result.project, 'cpython')
self.assertEqual(result.cve_id, 'cve_id')
self.assertIsNotNone(result.repository)
self.assertIsNotNone(result.references)
# custom parameters
prep = NVDFeedPreprocessor(
attributes=TEST_CVE_ATTR # only extract cve_id
)
result, = prep.fit_transform(X=[TEST_CVE])
self.assertEqual(result.user, 'python')
self.assertEqual(result.project, 'cpython')
self.assertEqual(result.cve_id, 'cve_id')
self.assertIsNotNone(result.repository)
self.assertIsNotNone(result.references)
# plus description and CVE_ATTR_VALS
self.assertEqual(result[-len(TEST_CVE_ATTR):], TEST_CVE_ATTR_VALS)
self.assertEqual(result.description, 'description')
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')