Skip to content

Commit

Permalink
Fix typing in telegram provider (#40255)
Browse files Browse the repository at this point in the history
The python-telegram-bot new version has typing added and we should
pass the right dict type to it.
  • Loading branch information
potiuk authored Jun 15, 2024
1 parent 35871f8 commit bc4ca9d
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 10 deletions.
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

0 comments on commit bc4ca9d

Please sign in to comment.