feat: gate post generation on training flag

This commit is contained in:
admin committed 2026-10-09 20:11:41 +02:00
1 parent e847a4b390
commit a43ededa1b
5 files changed
+132 -6

No files matched your search

+4 -4
View File
@@ -239,19 +239,19 @@ class ModeRoutingTests(unittest.IsolatedAsyncioTestCase):
async def test_post_command_generates_only_in_generation_mode(self):
message = _StubMessage("/post")
state = _StubState("generation")
with patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify:
with patch.object(bot, "_generate_post_from_training", new=AsyncMock()) as generate:
await bot.handle_command_text(message, state)
classify.assert_awaited_once_with(message, "Напиши пост", state=state)
generate.assert_awaited_once_with(message)
async def test_generate_post_button_runs_post_command(self):
message = _StubMessage(keyboards.POST_GENERATE_BUTTON)
state = _StubState("generation")
with (
patch.object(bot, "notify_development_deploy", new=AsyncMock()),
patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify,
patch.object(bot, "_generate_post_from_training", new=AsyncMock()) as generate,
):
await bot.handle_text(message, state)
classify.assert_awaited_once_with(message, "Напиши пост", state=state)
generate.assert_awaited_once_with(message)
async def test_training_text_is_saved_as_note_without_entering_generation(self):
message = _StubMessage("Кофейня в Калининграде для семей с детьми")
+46
View File
@@ -475,6 +475,48 @@ class GenerationClassificationRoutingTests(unittest.IsolatedAsyncioTestCase):
answer=AsyncMock(),
)
async def test_generate_post_refuses_when_training_flag_is_off(self):
message = self.make_message()
with (
patch.object(bot, "_check_subscription", new=AsyncMock(return_value=True)),
patch.object(bot, "is_user_trained", return_value=False),
patch.object(bot.asyncio, "to_thread", new=AsyncMock(return_value=False)),
patch.object(bot, "generate_posts", new=AsyncMock()) as generate,
):
await bot._generate_post_from_training(cast(Any, message))
generate.assert_not_awaited()
message.answer.assert_awaited_once()
self.assertIn("Сначала обучи меня", message.answer.await_args.args[0])
async def test_generate_post_uses_training_material_without_intent_classifier(self):
message = self.make_message()
status = AsyncMock()
async def run_in_thread(function, *args):
return function(*args)
with (
patch.object(bot, "_check_subscription", new=AsyncMock(return_value=True)),
patch.object(bot, "is_user_trained", return_value=True),
patch.object(bot.asyncio, "to_thread", side_effect=run_in_thread),
patch.object(bot, "get_tone_of_voice", return_value="Кофейня для семей"),
patch.object(bot, "get_post_samples", return_value=["Открыли террасу"]),
patch.object(bot, "get_training_notes", return_value=["Работаем в Калининграде"]),
patch.object(bot.OperationStatus, "start", new=AsyncMock(return_value=status)),
patch.object(bot, "generate_posts", new=AsyncMock(return_value=["A", "B", "C"])) as generate,
patch.object(bot, "_send_three_posts", new=AsyncMock()) as send_posts,
patch.object(bot, "classify_intent", new=AsyncMock()) as classify,
):
await bot._generate_post_from_training(cast(Any, message))
classify.assert_not_awaited()
self.assertIsNotNone(generate.await_args)
args = generate.await_args.args
self.assertEqual(args[0]["intent"], "write_post")
self.assertIn("Сам выбери одну конкретную", args[0]["requirements"])
self.assertIn("Открыли террасу", args[0]["requirements"])
self.assertIn("Работаем в Калининграде", args[0]["requirements"])
send_posts.assert_awaited_once_with(message, ["A", "B", "C"], user_id=91, origin="generated")
async def test_thanks_and_other_use_prepared_replies_without_generators(self):
for label, intent, replies in (
("thanks", "thanks", bot.THANKS_REPLIES),
@@ -682,11 +724,14 @@ class ScreenshotToneWorkflowTests(unittest.IsolatedAsyncioTestCase):
with tempfile.TemporaryDirectory() as temp_dir:
db = str(Path(temp_dir) / "training.sqlite3")
user_store.init_db(db)
self.assertFalse(user_store.is_user_trained(23, db))
self.assertEqual(user_store.add_training_note(23, " Кофейня в Калининграде ", db), ["Кофейня в Калининграде"])
self.assertTrue(user_store.is_user_trained(23, db))
user_store.add_post_samples(23, ["Пост об открытии"], db)
self.assertEqual(user_store.get_training_notes(23, db), ["Кофейня в Калининграде"])
self.assertEqual(user_store.get_post_samples(23, db), ["Пост об открытии"])
user_store.clear_user_data(23, db)
self.assertFalse(user_store.is_user_trained(23, db))
self.assertEqual(user_store.get_training_notes(23, db), [])
self.assertEqual(user_store.get_post_samples(23, db), [])
@@ -852,6 +897,7 @@ class ScreenshotToneWorkflowTests(unittest.IsolatedAsyncioTestCase):
connection.execute("INSERT INTO users (telegram_user_id, tone_of_voice) VALUES (1, 'старый стиль')")
user_store.init_db(db)
self.assertEqual(user_store.get_tone_of_voice(1, db), "старый стиль")
self.assertTrue(user_store.is_user_trained(1, db))
self.assertEqual(user_store.add_post_samples(1, ["Сохранённый пост"], db), ["Сохранённый пост"])