From a55baa613cbdc7eac6624b2030fa9c9f9d815309 Mon Sep 17 00:00:00 2001
From: emmatveev <emmtvv@gmail.com>
Date: Wed, 27 Nov 2024 02:48:30 +0300
Subject: [PATCH] fix

---
 app/routers/take.py        | 4 ++--
 app/storage/mongo/tasks.py | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/app/routers/take.py b/app/routers/take.py
index b56bd1a..e115551 100644
--- a/app/routers/take.py
+++ b/app/routers/take.py
@@ -25,5 +25,5 @@ class Response(pydantic.BaseModel):
 async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
     task = await tasks.take_task(queue)
     if not task:
-        raise fastapi.HTTPException(404)
-    return Response(id=str(task._id), attempt=task.attempts, payload=task.payload, task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))
+        return Response(task=None)
+    return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))
diff --git a/app/storage/mongo/tasks.py b/app/storage/mongo/tasks.py
index e1b480a..ec98f36 100644
--- a/app/storage/mongo/tasks.py
+++ b/app/storage/mongo/tasks.py
@@ -1,7 +1,6 @@
 import bson
 import datetime
 import pydantic
-import typing
 
 from app.storage.mongo import database
 from app.utils import time
@@ -27,7 +26,7 @@ async def add_task(task: Task) -> str:
     return result.inserted_id
 
 
-async def take_task(queue: str) -> typing.Optional[Task]:
+async def take_task(queue: str) -> Task|None:
     now = time.now()
     async for raw_task in collection.find({'queue': queue, 'available_from': {'$lte': now}}):
         task = Task.model_validate(raw_task)