langchain/tests/unit_tests/document_loaders/test_telegram.py
Raduan Al-Shedivat 00c6ec8a2d
fix(document_loaders/telegram): fix pandas calls + add tests (#4806)
# Fix Telegram API loader + add tests.
I was testing this integration and it was broken with next error:
```python
message_threads = loader._get_message_threads(df)
KeyError: False
```
Also, this particular loader didn't have any tests / related group in
poetry, so I added those as well.

@hwchase17 / @eyurtsev please take a look on this fix PR.

---------

Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
2023-05-16 14:35:25 -07:00

37 lines
1.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from pathlib import Path
import pytest
from langchain.document_loaders import TelegramChatApiLoader, TelegramChatFileLoader
def test_telegram_chat_file_loader() -> None:
"""Test TelegramChatFileLoader."""
file_path = Path(__file__).parent / "test_docs" / "telegram.json"
loader = TelegramChatFileLoader(str(file_path))
docs = loader.load()
assert len(docs) == 1
assert docs[0].metadata["source"] == str(file_path)
assert docs[0].page_content == (
"Henry on 2020-01-01T00:00:02: It's 2020...\n\n"
"Henry on 2020-01-01T00:00:04: Fireworks!\n\n"
"Grace 🧤 ðŸ\x8d on 2020-01-01T00:00:05: You're a minute late!\n\n"
)
@pytest.mark.requires("telethon", "pandas")
def test_telegram_channel_loader_parsing() -> None:
"""Test TelegramChatApiLoader."""
file_path = Path(__file__).parent / "test_docs" / "telegram_channel.json"
# if we don't provide any value, it will skip fetching from telegram
# and will check the parsing logic.
loader = TelegramChatApiLoader(file_path=str(file_path))
docs = loader.load()
assert len(docs) == 1
print(docs[0].page_content)
assert docs[0].page_content == (
"Hello, world!.\nLLMs are awesome! Langchain is great. Telegram is the best!."
)