File size: 1,837 Bytes
854f61d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from os import name
import unittest
from src.knowledge_service.knowledge_retrieval import *
import pandas as pd
class InformationRetrievalTests(unittest.TestCase):
def test_get_candidates_with_standard_name_retrieves_ok(self):
input = 'A Voice for Men'
index = 'dangerous_organizations'
results = get_candidates(input, index)
self.assertIsNotNone(results)
expected = {'name': 'A Voice for Men', 'summary': 'Organization', 'label': ['ORGANIZATION']}
self.assertEqual(expected, results[0])
def test_get_candidates_with_alternate_latin_name_retrieves_ok(self):
input = 'Uma Voz'
index = 'dangerous_organizations'
results = get_candidates(input, index)
self.assertIsNotNone(results)
expected = {'name': 'A Voice for Men', 'summary': 'Organization', 'label': ['ORGANIZATION']}
self.assertEqual(expected, results[0])
def test_get_information_standard_individual_returns_ok(self):
entity = 'Curt Doolittle'
index = 'dangerous_individuals'
result = get_information(entity=entity, index=index)
print(result)
self.assertIsNotNone(result)
self.assertIn("""Name: Curt Doolittle
Gender: Male
Summary: CEO at A. O. Smith
Policy Category: HATE""", result)
def test_get_information_standard_organization_returns_ok(self):
entity = 'Voice for Men'
index = 'dangerous_organizations'
result = get_information(entity=entity, index=index)
print(result)
self.assertIsNotNone(result)
self.assertIn("Name: A Voice for Men", result)
if __name__=='__main__':
unittest.main()
|