Refactor Telegram tests

joweich · Dec 3, 2023 · 030c7b3 · 030c7b3
1 parent 1e673a7
commit 030c7b3
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 25 deletions.
diff --git a/test/telegram/target.json b/test/telegram/target.json
@@ -0,0 +1,12 @@
+[
+    {
+        "timestamp": "2018-12-05T11:39:10",
+        "author": "Author 1",
+        "message": "Minimal Message"
+    },
+    {
+        "timestamp": "2018-12-05T11:49:24",
+        "author": "Author 2",
+        "message": "Link: https://telegram.org/"
+    }
+]
diff --git a/test/telegram/test_target.csv b/test/telegram/test_target.csv
diff --git a/test/test_telegram.py b/test/test_telegram.py
@@ -1,32 +1,17 @@
-import pandas as pd
-from pandas.testing import assert_frame_equal
-
-from chatminer.chatparsers import TelegramJsonParser
+from chatminer.chatparsers import ParsedMessageCollection, TelegramJsonParser
 
 
 def test_telegram_single_export():
+    target = ParsedMessageCollection()
+    target.read_from_json("test/telegram/target.json")
     parser = TelegramJsonParser("test/telegram/test_single_export.json")
     parser.parse_file()
-    df_res = parser.parsed_messages.get_df()
-    df_test = pd.read_csv(
-        "test/telegram/test_target.csv",
-        parse_dates=["timestamp"],
-    )
-    assert_frame_equal(
-        df_test[["author", "message", "words", "letters"]],
-        df_res[["author", "message", "words", "letters"]],
-    )
+    assert parser.parsed_messages == target
 
 
 def test_telegram_batch_export():
+    target = ParsedMessageCollection()
+    target.read_from_json("test/telegram/target.json")
     parser = TelegramJsonParser("test/telegram/test_batch_export.json", "Chatname")
     parser.parse_file()
-    df_res = parser.parsed_messages.get_df()
-    df_test = pd.read_csv(
-        "test/telegram/test_target.csv",
-        parse_dates=["timestamp"],
-    )
-    assert_frame_equal(
-        df_test[["author", "message", "words", "letters"]],
-        df_res[["author", "message", "words", "letters"]],
-    )
+    assert parser.parsed_messages == target