Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix typing in telegram provider #40255

Merged
merged 1 commit into from
Jun 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions airflow/providers/telegram/hooks/telegram.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,17 +142,18 @@ def send_message(self, api_params: dict) -> None:

:param api_params: params for telegram_instance.send_message. It can also be used to override chat_id
"""
kwargs = {
"chat_id": self.chat_id,
kwargs: dict[str, int | str | bool] = {
"parse_mode": telegram.constants.ParseMode.HTML,
"disable_web_page_preview": True,
}
if self.chat_id is not None:
kwargs["chat_id"] = self.chat_id
kwargs.update(api_params)

if "text" not in kwargs or kwargs["text"] is None:
raise AirflowException("'text' must be provided for telegram message")

if kwargs["chat_id"] is None:
if kwargs.get("chat_id") is None:
raise AirflowException("'chat_id' must be provided for telegram message")

response = asyncio.run(self.connection.send_message(**kwargs))
Expand Down
14 changes: 7 additions & 7 deletions tests/providers/telegram/hooks/test_telegram.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,21 +93,21 @@ def test_should_raise_exception_if_message_text_is_not_provided(self, mock_get_c
hook = TelegramHook(telegram_conn_id="telegram_default")
error_message = "'text' must be provided for telegram message"
with pytest.raises(airflow.exceptions.AirflowException, match=error_message):
hook.send_message({"chat_id": -420913222})
hook.send_message({"chat_id": "-420913222"})

@mock.patch("airflow.providers.telegram.hooks.telegram.TelegramHook.get_conn")
def test_should_send_message_if_all_parameters_are_correctly_provided(self, mock_get_conn):
mock_get_conn.return_value = AsyncMock(password="some_token")

hook = TelegramHook(telegram_conn_id="telegram_default")
hook.send_message({"chat_id": -420913222, "text": "test telegram message"})
hook.send_message({"chat_id": "-420913222", "text": "test telegram message"})

mock_get_conn.return_value.send_message.return_value = "OK."

mock_get_conn.assert_called_once()
mock_get_conn.return_value.send_message.assert_called_once_with(
**{
"chat_id": -420913222,
"chat_id": "-420913222",
"parse_mode": "HTML",
"disable_web_page_preview": True,
"text": "test telegram message",
Expand All @@ -118,15 +118,15 @@ def test_should_send_message_if_all_parameters_are_correctly_provided(self, mock
def test_should_send_message_if_chat_id_is_provided_through_constructor(self, mock_get_conn):
mock_get_conn.return_value = AsyncMock(password="some_token")

hook = TelegramHook(telegram_conn_id="telegram_default", chat_id=-420913222)
hook = TelegramHook(telegram_conn_id="telegram_default", chat_id="-420913222")
hook.send_message({"text": "test telegram message"})

mock_get_conn.return_value.send_message.return_value = "OK."

mock_get_conn.assert_called_once()
mock_get_conn.return_value.send_message.assert_called_once_with(
**{
"chat_id": -420913222,
"chat_id": "-420913222",
"parse_mode": "HTML",
"disable_web_page_preview": True,
"text": "test telegram message",
Expand Down Expand Up @@ -182,15 +182,15 @@ def side_effect(*args, **kwargs):
def test_should_send_message_if_token_is_provided(self, mock_get_conn):
mock_get_conn.return_value = AsyncMock(password="some_token")

hook = TelegramHook(token=TELEGRAM_TOKEN, chat_id=-420913222)
hook = TelegramHook(token=TELEGRAM_TOKEN, chat_id="-420913222")
hook.send_message({"text": "test telegram message"})

mock_get_conn.return_value.send_message.return_value = "OK."

mock_get_conn.assert_called_once()
mock_get_conn.return_value.send_message.assert_called_once_with(
**{
"chat_id": -420913222,
"chat_id": "-420913222",
"parse_mode": "HTML",
"disable_web_page_preview": True,
"text": "test telegram message",
Expand Down
Loading