import unittest

import sys, os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

from digExtractor.extractor_processor import ExtractorProcessor
from digCveExtractor.cve_extractor import CveExtractor

class TestCveExtractorMethods(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_cve_extractor(self):
        doc = {'text' : 'Sample cves are CVE-1993-1344 ascasdsadfjskdvnjdvn CVE-2006-1232'}

        extractor = CveExtractor().set_metadata({'extractor': 'cve'})
        ep = ExtractorProcessor().set_input_fields(['text'])\
                                 .set_output_field('extracted')\
                                 .set_extractor(extractor)
        updated_doc = ep.extract(doc)
        result = updated_doc['extracted'][0]['result']
        self.assertEqual(result[0]['value'],
                         'CVE-1993-1344')
        self.assertEqual(result[1]['value'],
                         'CVE-2006-1232')

if __name__ == '__main__':
    unittest.main()
