From 0d30e93dd1b1c6b342ebae06de478a6c099ed435 Mon Sep 17 00:00:00 2001
From: Desiders <abc_desi@proton.me>
Date: Thu, 13 Jun 2024 23:46:26 +0300
Subject: [PATCH] Parametrize test

---
 .../test_middlewares/test_user_context.py     | 49 +++++++------------
 1 file changed, 19 insertions(+), 30 deletions(-)

diff --git a/tests/test_dispatcher/test_middlewares/test_user_context.py b/tests/test_dispatcher/test_middlewares/test_user_context.py
index 6ba9c07b..3fc87eac 100644
--- a/tests/test_dispatcher/test_middlewares/test_user_context.py
+++ b/tests/test_dispatcher/test_middlewares/test_user_context.py
@@ -15,6 +15,7 @@ from aiogram.types import (
     ChatBoost,
     ChatBoostSourcePremium,
     ChatBoostSourceGiftCode,
+    ChatBoostSourceGiveaway,
 )
 
 
@@ -51,16 +52,30 @@ class TestUserContextMiddleware:
         assert data["event_from_user"] is user
         assert data["event_thread_id"] == thread_id
 
-    async def test_resolve_event_context(self):
+    @pytest.mark.parametrize(
+        "source, expected_user",
+        [
+            (
+                ChatBoostSourcePremium(user=User(id=2, first_name="Test", is_bot=False)),
+                User(id=2, first_name="Test", is_bot=False),
+            ),
+            (ChatBoostSourceGiftCode(user=User(id=2, first_name="Test", is_bot=False)), None),
+            (
+                ChatBoostSourceGiveaway(
+                    giveaway_message_id=1, user=User(id=2, first_name="Test", is_bot=False)
+                ),
+                None,
+            ),
+        ],
+    )
+    async def test_resolve_event_context(self, source, expected_user):
         middleware = UserContextMiddleware()
         data = {}
 
         chat = Chat(id=1, type="private", title="Test")
-        user = User(id=2, first_name="Test", is_bot=False)
         add_date = datetime.now()
         expiration_date = datetime.now()
 
-        source = ChatBoostSourcePremium(user=user)
         boost = ChatBoost(
             boost_id="Test", add_date=add_date, expiration_date=expiration_date, source=source
         )
@@ -70,30 +85,4 @@ class TestUserContextMiddleware:
 
         event_context = data["event_context"]
         assert isinstance(event_context, EventContext)
-        assert event_context.chat is chat
-        assert event_context.user is user
-        assert event_context.thread_id is None
-        assert data["event_chat"] is chat
-        assert data["event_from_user"] is user
-        assert data.get("event_thread_id") is None
-
-        data.clear()
-
-        source = ChatBoostSourceGiftCode(user=user)
-        boost = ChatBoost(
-            boost_id="Test", add_date=add_date, expiration_date=expiration_date, source=source
-        )
-        update = Update(update_id=42, chat_boost=ChatBoostUpdated(chat=chat, boost=boost))
-
-        await middleware(next_handler, update, data)
-
-        event_context = data["event_context"]
-        assert isinstance(event_context, EventContext)
-        assert event_context.chat is chat
-        assert event_context.user is None
-        assert event_context.thread_id is None
-        assert data["event_chat"] is chat
-        assert data.get("event_from_user") is None
-        assert data.get("event_thread_id") is None
-
-        data.clear()
+        assert event_context.user == expected_user