# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#
import unittest
from typing import Any, Dict
from google.cloud.language_v1.proto.language_service_pb2 import Document
from airflow.providers.google.cloud.hooks.natural_language import CloudNaturalLanguageHook
from tests.compat import mock
from tests.gcp.utils.base_gcp_mock import mock_base_gcp_hook_no_default_project_id
API_RESPONSE = {} # type: Dict[Any, Any]
DOCUMENT = Document(
content="Airflow is a platform to programmatically author, schedule and monitor workflows."
)
ENCODING_TYPE = "UTF32"
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
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.natural_language import (
CloudNaturalLanguageAnalyzeEntitiesOperator, CloudNaturalLanguageAnalyzeEntitySentimentOperator,
CloudNaturalLanguageAnalyzeSentimentOperator, CloudNaturalLanguageClassifyTextOperator,
)
from airflow.utils.dates import days_ago
# [START howto_operator_gcp_natural_language_document_text]
TEXT = """Airflow is a platform to programmatically author, schedule and monitor workflows.
Use Airflow to author workflows as Directed Acyclic Graphs (DAGs) of tasks. The Airflow scheduler executes
your tasks on an array of workers while following the specified dependencies. Rich command line utilities
make performing complex surgeries on DAGs a snap. The rich user interface makes it easy to visualize
pipelines running in production, monitor progress, and troubleshoot issues when needed.
"""
document = Document(content=TEXT, type="PLAIN_TEXT")
# [END howto_operator_gcp_natural_language_document_text]
# [START howto_operator_gcp_natural_language_document_gcs]
GCS_CONTENT_URI = "gs://my-text-bucket/sentiment-me.txt"
document_gcs = Document(gcs_content_uri=GCS_CONTENT_URI, type="PLAIN_TEXT")
# [END howto_operator_gcp_natural_language_document_gcs]
default_args = {"start_date": days_ago(1)}
with models.DAG(
"example_gcp_natural_language",
default_args=default_args,
schedule_interval=None, # Override to match your needs
) as dag:
