Ejemplo n.º 1
0
    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)
    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(  # type: ignore
        "airflow.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(result, API_RESPONSE)

        get_conn.return_value.classify_text.assert_called_once_with(
            document=DOCUMENT, retry=None, timeout=None, metadata=None
        )
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(  # type: ignore
        "airflow.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(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.contrib.hooks.gcp_natural_language_hook.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.assertIs(result, API_RESPONSE)

        get_conn.return_value.classify_text.assert_called_once_with(
            document=DOCUMENT, retry=None, timeout=None, metadata=None)