fix task hooks

Adori · Dec 16, 2021 · 8146205 · 8146205
1 parent cc47cee
commit 8146205
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/fastapi_cloud_tasks/hooks.py b/fastapi_cloud_tasks/hooks.py
@@ -23,7 +23,7 @@ def chain(request):
     return chain
 
 
-def oidc_scheduler_hook(token: tasks_v2.OidcToken) -> SchedulerHook:
+def oidc_scheduler_hook(token: scheduler_v1.OidcToken) -> SchedulerHook:
     def add_token(
         request: scheduler_v1.CreateJobRequest,
     ) -> scheduler_v1.CreateJobRequest:
@@ -33,7 +33,7 @@ def add_token(
     return add_token
 
 
-def oidc_task_hook(token: scheduler_v1.OidcToken) -> TaskHook:
+def oidc_task_hook(token: tasks_v2.OidcToken) -> TaskHook:
     def add_token(request: tasks_v2.CreateTaskRequest) -> tasks_v2.CreateTaskRequest:
         request.task.http_request.oidc_token = token
         return request