def test_process(meta: Dict[str, Any], expected: List[Dict[str, Any]]) -> None:
training_data = mock.MagicMock()
message = mock.MagicMock()
message.data = {"text": 'input "test"'}
message.get.return_value = []
training_data.regex_features = []
extractor = RegexEntityExtractor(meta)
extractor.process(message)
message.set.assert_called_once_with("entities", expected, add_to_output=True)
def test_persist(
write_json_to_file: mock.MagicMock, regex_features: Dict[str, Any], expected: Dict[str, Any], persisted: bool,
) -> None:
training_data = mock.MagicMock()
training_data.regex_features = regex_features
extractor = RegexEntityExtractor({})
extractor.train(training_data)
assert expected == extractor.persist("filename", "model_dir")
if persisted:
write_json_to_file.assert_called_once_with(os.path.join("model_dir", "filename.json"), regex_features)
else:
write_json_to_file.assert_not_called()
def test_init(re_compile: mock.MagicMock, meta: Dict[str, Any], expected: List[Dict[str, Any]]) -> None:
extractor = RegexEntityExtractor(meta)
assert extractor.regex_features == expected
if len(expected):
re_compile.assert_called_once()
else:
re_compile.assert_not_called()
def test_load(read_json_file: mock.MagicMock, meta: Dict[str, Any], loaded: bool) -> None:
extractor = RegexEntityExtractor.load(meta, "model_dir")
if loaded:
read_json_file.assert_called_once_with(os.path.join("model_dir", "test.json"))
assert isinstance(extractor, RegexEntityExtractor)
else:
read_json_file.assert_not_called()
def test_train() -> None:
training_data = mock.MagicMock()
training_data.regex_features = []
extractor = RegexEntityExtractor({})
extractor.train(training_data)
assert extractor.regex_features == training_data.regex_features