def setUp(self):
with mock.patch(
"airflow.contrib.hooks."
"gcp_api_base_hook.CloudBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = CloudNaturalLanguageHook(gcp_conn_id="test")
def execute(self, context):
hook = CloudNaturalLanguageHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Start analyzing entities")
response = hook.analyze_entities(document=self.document,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata)
self.log.info("Finished analyzing entities")
return MessageToDict(response)
def execute(self, context):
hook = CloudNaturalLanguageHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Start text classify")
response = hook.classify_text(document=self.document,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata)
self.log.info("Finished text classify")
return MessageToDict(response)
def execute(self, context):
hook = CloudNaturalLanguageHook(
gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain,
)
self.log.info("Start sentiment analyze")
response = hook.analyze_sentiment(
document=self.document, retry=self.retry, timeout=self.timeout, metadata=self.metadata
)
self.log.info("Finished sentiment analyze")
return MessageToDict(response)
def execute(self, context):
hook = CloudNaturalLanguageHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Start entity sentiment analyze")
response = hook.analyze_entity_sentiment(
document=self.document,
encoding_type=self.encoding_type,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
self.log.info("Finished entity sentiment analyze")
return MessageToDict(response)
class TestCloudNaturalLanguageHook(unittest.TestCase):
def setUp(self):
with mock.patch(
"airflow.contrib.hooks." "gcp_api_base_hook.GoogleCloudBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = CloudNaturalLanguageHook(gcp_conn_id="test")
@mock.patch(
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.client_info",
new_callable=mock.PropertyMock
)
@mock.patch("airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook."
"_get_credentials")
@mock.patch("airflow.providers.google.cloud.hooks.natural_language.LanguageServiceClient")
def test_language_service_client_creation(self, mock_client, mock_get_creds, mock_client_info):
result = self.hook.get_conn()
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value
)
self.assertEqual(mock_client.return_value, result)
self.assertEqual(self.hook._conn, result)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.get_conn",
**{"return_value.analyze_entities.return_value": API_RESPONSE} # type: ignore
)
def test_analyze_entities(self, get_conn):
result = self.hook.analyze_entities(document=DOCUMENT, encoding_type=ENCODING_TYPE)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.analyze_entities.assert_called_once_with(
document=DOCUMENT, encoding_type=ENCODING_TYPE, retry=None, timeout=None, metadata=None
)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.get_conn",
**{"return_value.analyze_entity_sentiment.return_value": API_RESPONSE}
)
def test_analyze_entity_sentiment(self, get_conn):
result = self.hook.analyze_entity_sentiment(document=DOCUMENT, encoding_type=ENCODING_TYPE)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.analyze_entity_sentiment.assert_called_once_with(
document=DOCUMENT, encoding_type=ENCODING_TYPE, retry=None, timeout=None, metadata=None
)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.get_conn",
**{"return_value.analyze_sentiment.return_value": API_RESPONSE}
)
def test_analyze_sentiment(self, get_conn):
result = self.hook.analyze_sentiment(document=DOCUMENT, encoding_type=ENCODING_TYPE)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.analyze_sentiment.assert_called_once_with(
document=DOCUMENT, encoding_type=ENCODING_TYPE, retry=None, timeout=None, metadata=None
)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.get_conn",
**{"return_value.analyze_syntax.return_value": API_RESPONSE}
)
def test_analyze_syntax(self, get_conn):
result = self.hook.analyze_syntax(document=DOCUMENT, encoding_type=ENCODING_TYPE)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.analyze_syntax.assert_called_once_with(
document=DOCUMENT, encoding_type=ENCODING_TYPE, retry=None, timeout=None, metadata=None
)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.get_conn",
**{"return_value.annotate_text.return_value": API_RESPONSE}
)
def test_annotate_text(self, get_conn):
result = self.hook.annotate_text(document=DOCUMENT, encoding_type=ENCODING_TYPE, features=None)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.annotate_text.assert_called_once_with(
document=DOCUMENT,
encoding_type=ENCODING_TYPE,
features=None,
retry=None,
timeout=None,
metadata=None,
)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.natural_language.CloudNaturalLanguageHook.get_conn",
**{"return_value.classify_text.return_value": API_RESPONSE}
)
def test_classify_text(self, get_conn):
result = self.hook.classify_text(document=DOCUMENT)
self.assertEqual(result, API_RESPONSE)
get_conn.return_value.classify_text.assert_called_once_with(
document=DOCUMENT, retry=None, timeout=None, metadata=None
)
def setUp(self):
with mock.patch(
"airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = CloudNaturalLanguageHook(gcp_conn_id="test")
