test_message_cleanup.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  1. from __future__ import annotations
  2. import asyncio
  3. from types import SimpleNamespace
  4. import pytest
  5. class FakeMessage:
  6. def __init__(self, chat_id: int, message_id: int, *, text: str = "") -> None:
  7. self.chat = SimpleNamespace(
  8. id=chat_id,
  9. type=SimpleNamespace(value="supergroup" if chat_id < 0 else "private"),
  10. )
  11. self.id = message_id
  12. self.text = text
  13. self.caption = None
  14. self.delete_calls: list[bool] = []
  15. async def delete(self, *, revoke: bool = True) -> None:
  16. self.delete_calls.append(revoke)
  17. class FakeClient:
  18. def __init__(self) -> None:
  19. self.next_message_id = 100
  20. async def send_message(
  21. self,
  22. chat_id: int,
  23. _text: str,
  24. *,
  25. reply_to_message_id: int | None = None,
  26. ) -> FakeMessage:
  27. self.next_message_id += 1
  28. return FakeMessage(chat_id, self.next_message_id)
  29. @pytest.fixture
  30. def cleanup(app_modules):
  31. module = app_modules.load("wbb.services.message_cleanup")
  32. module._scheduled_messages.clear()
  33. module._cleanup_tasks.clear()
  34. module._active_interactions.clear()
  35. yield module
  36. module._scheduled_messages.clear()
  37. module._cleanup_tasks.clear()
  38. module._active_interactions.clear()
  39. async def test_group_cleanup_revokes_once_for_everyone(cleanup, monkeypatch):
  40. real_sleep = asyncio.sleep
  41. async def immediate_sleep(_delay: int) -> None:
  42. return None
  43. monkeypatch.setattr(cleanup.asyncio, "sleep", immediate_sleep)
  44. message = FakeMessage(-1001, 10, text="/points")
  45. cleanup.schedule_group_cleanup(message)
  46. cleanup.schedule_group_cleanup(message)
  47. await real_sleep(0)
  48. await real_sleep(0)
  49. assert message.delete_calls == [True]
  50. assert (-1001, 10) not in cleanup._scheduled_messages
  51. async def test_private_messages_are_not_scheduled(cleanup):
  52. message = FakeMessage(1001, 11, text="/points")
  53. cleanup.schedule_group_cleanup(message)
  54. assert not cleanup.is_group_cleanup_scheduled(message)
  55. assert message.delete_calls == []
  56. async def test_cleanup_task_is_referenced_until_deletion_finishes(
  57. cleanup, monkeypatch
  58. ):
  59. real_sleep = asyncio.sleep
  60. release = asyncio.Event()
  61. async def blocked_sleep(_delay: int) -> None:
  62. await release.wait()
  63. monkeypatch.setattr(cleanup.asyncio, "sleep", blocked_sleep)
  64. message = FakeMessage(-1001, 17, text="签到")
  65. cleanup.schedule_group_cleanup(message)
  66. await real_sleep(0)
  67. assert len(cleanup._cleanup_tasks) == 1
  68. release.set()
  69. await real_sleep(0)
  70. await real_sleep(0)
  71. assert message.delete_calls == [True]
  72. assert cleanup._cleanup_tasks == set()
  73. async def test_command_context_tracks_only_responses_in_the_same_group(
  74. cleanup, monkeypatch
  75. ):
  76. scheduled: list[FakeMessage] = []
  77. monkeypatch.setattr(
  78. cleanup,
  79. "schedule_group_cleanup",
  80. lambda *messages, **_kwargs: scheduled.extend(messages),
  81. )
  82. client = FakeClient()
  83. cleanup.install_group_interaction_tracking(client)
  84. command = FakeMessage(-1001, 12, text="/points")
  85. assert cleanup.begin_group_message_interaction(command)
  86. group_response = await client.send_message(
  87. -1001, "积分余额", reply_to_message_id=command.id
  88. )
  89. await client.send_message(-1001, "抽奖主消息")
  90. await client.send_message(-1002, "其他群通知", reply_to_message_id=command.id)
  91. await client.send_message(1001, "私聊通知", reply_to_message_id=command.id)
  92. assert scheduled == [command, group_response]
  93. async def test_plain_message_clears_previous_command_context(cleanup, monkeypatch):
  94. scheduled: list[FakeMessage] = []
  95. monkeypatch.setattr(
  96. cleanup,
  97. "schedule_group_cleanup",
  98. lambda *messages, **_kwargs: scheduled.extend(messages),
  99. )
  100. client = FakeClient()
  101. cleanup.install_group_interaction_tracking(client)
  102. command = FakeMessage(-1001, 13, text="/points")
  103. plain_message = FakeMessage(-1001, 14, text="大家好")
  104. assert cleanup.begin_group_message_interaction(command)
  105. assert not cleanup.begin_group_message_interaction(plain_message)
  106. await client.send_message(
  107. -1001, "普通自动回复", reply_to_message_id=plain_message.id
  108. )
  109. assert scheduled == [command]
  110. async def test_checkin_button_is_treated_as_group_interaction(cleanup, monkeypatch):
  111. scheduled: list[FakeMessage] = []
  112. monkeypatch.setattr(
  113. cleanup,
  114. "schedule_group_cleanup",
  115. lambda *messages, **_kwargs: scheduled.extend(messages),
  116. )
  117. message = FakeMessage(-1001, 15, text="签到")
  118. assert cleanup.begin_group_message_interaction(message, menu_texts={"签到"})
  119. assert scheduled == [message]
  120. async def test_callback_only_continues_an_already_ephemeral_menu(
  121. cleanup, monkeypatch
  122. ):
  123. scheduled: list[FakeMessage] = []
  124. monkeypatch.setattr(
  125. cleanup,
  126. "schedule_group_cleanup",
  127. lambda *messages, **_kwargs: scheduled.extend(messages),
  128. )
  129. client = FakeClient()
  130. cleanup.install_group_interaction_tracking(client)
  131. persistent_menu = FakeMessage(-1001, 16)
  132. query = SimpleNamespace(message=persistent_menu)
  133. assert not cleanup.begin_group_callback_interaction(query)
  134. await client.send_message(
  135. -1001, "不应自动删除", reply_to_message_id=persistent_menu.id
  136. )
  137. cleanup._scheduled_messages.add((-1001, 16))
  138. assert cleanup.begin_group_callback_interaction(query)
  139. callback_response = await client.send_message(
  140. -1001, "临时菜单回复", reply_to_message_id=persistent_menu.id
  141. )
  142. assert scheduled == [callback_response]