follow-resolver-fix2

resolvers/follower.py

@@ -23,6 +23,17 @@ from storage.schema import mutation, query
 from utils.logger import root_logger as logger
 
 
+def get_entity_field_name(entity_type: str) -> str:
+    """Возвращает имя поля для связи с сущностью в модели подписчика"""
+    entity_field_mapping = {
+        "author": "following",
+        "topic": "topic", 
+        "community": "community",
+        "shout": "shout"
+    }
+    return entity_field_mapping[entity_type]
+
+
 @mutation.field("follow")
 @login_required
 async def follow(
@@ -88,12 +99,13 @@ async def follow(
             logger.debug(f"entity_id: {entity_id}, entity_dict: {entity_dict}")
 
             if entity_id is not None and isinstance(entity_id, int):
+                entity_field = get_entity_field_name(entity_type)
+                
                 existing_sub = (
                     session.query(follower_class)
                     .where(
                         follower_class.follower == follower_id,  # type: ignore[attr-defined]
-                        follower_class.following == entity_id,  # type: ignore[attr-defined]
-                        getattr(follower_class, entity_type) == entity_id,  # type: ignore[attr-defined]
+                        getattr(follower_class, entity_field) == entity_id,  # type: ignore[attr-defined]
                     )
                     .first()
                 )
@@ -102,7 +114,7 @@ async def follow(
                     error = "already following"
                 else:
                     logger.debug("Добавление новой записи в базу данных")
-                    sub = follower_class(follower=follower_id, **{entity_type: entity_id})
+                    sub = follower_class(follower=follower_id, **{entity_field: entity_id})
                     logger.debug(f"Создан объект подписки: {sub}")
                     session.add(sub)
                     session.commit()
@@ -210,14 +222,14 @@ async def unfollow(
                     return {"error": f"Cannot get ID for {what.lower()}"}
 
             logger.debug(f"entity_id: {entity_id}")
+            entity_field = get_entity_field_name(entity_type)
+            
             sub = (
                 session.query(follower_class)
                 .where(
                     and_(
                         follower_class.follower == follower_id,  # type: ignore[attr-defined]
-                        follower_class.following == entity_id
-                        if entity_type == "AUTHOR"
-                        else follower_class.topic == entity_id,  # type: ignore[attr-defined]
+                        getattr(follower_class, entity_field) == entity_id,  # type: ignore[attr-defined]
                     )
                 )
                 .first()