import unittest
from ner import NamedEntityRecognizer

class TestNamedEntityRecognizer(unittest.TestCase):
  def test_recognize(self):
    # Initialize the NamedEntityRecognizer
    model_name = "bert-base-uncased"
    ner = NamedEntityRecognizer(model_name)

    # Test input text
    text = "Apple Inc. is a technology company based in Cupertino, California."

    # Expected output
    expected_output = [
      {'entity': 'B-ORG', 'score': 0.999, 'index': 1, 'start': 0, 'end': 5, 'text': 'Apple'},
      {'entity': 'I-ORG', 'score': 0.999, 'index': 2, 'start': 6, 'end': 9, 'text': 'Inc.'},
      {'entity': 'O', 'score': 0.999, 'index': 3, 'start': 10, 'end': 11, 'text': '.'},
      {'entity': 'O', 'score': 0.999, 'index': 4, 'start': 12, 'end': 14, 'text': 'is'},
      {'entity': 'O', 'score': 0.999, 'index': 5, 'start': 15, 'end': 16, 'text': 'a'},
      {'entity': 'O', 'score': 0.999, 'index': 6, 'start': 17, 'end': 28, 'text': 'technology'},
      {'entity': 'O', 'score': 0.999, 'index': 7, 'start': 29, 'end': 37, 'text': 'company'},
      {'entity': 'O', 'score': 0.999, 'index': 8, 'start': 38, 'end': 43, 'text': 'based'},
      {'entity': 'O', 'score': 0.999, 'index': 9, 'start': 44, 'end': 46, 'text': 'in'},
      {'entity': 'B-LOC', 'score': 0.999, 'index': 10, 'start': 47, 'end': 56, 'text': 'Cupertino'},
      {'entity': 'I-LOC', 'score': 0.999, 'index': 11, 'start': 57, 'end': 67, 'text': 'California'},
      {'entity': 'O', 'score': 0.999, 'index': 12, 'start': 67, 'end': 68, 'text': '.'}
    ]

    # Call the recognize method
    output = ner.recognize(text)

    # Assert the output matches the expected output
    self.assertEqual(output, expected_output)

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