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