36
36
from celery .backends .database import DatabaseBackend
37
37
from celery .contrib .testing .worker import start_worker
38
38
from kombu .asynchronous import set_event_loop
39
+ from kubernetes .client import models as k8s
39
40
40
41
from airflow .configuration import conf
41
42
from airflow .exceptions import AirflowException , AirflowTaskTimeout
@@ -137,7 +138,37 @@ def _change_state(self, key: TaskInstanceKey, state: TaskInstanceState, info=Non
137
138
138
139
@pytest .mark .flaky (reruns = 3 )
139
140
@pytest .mark .parametrize ("broker_url" , _prepare_test_bodies ())
140
- def test_celery_integration (self , broker_url ):
141
+ @pytest .mark .parametrize (
142
+ "executor_config" ,
143
+ [
144
+ pytest .param ({}, id = "no_executor_config" ),
145
+ pytest .param (
146
+ {
147
+ "pod_override" : k8s .V1Pod (
148
+ spec = k8s .V1PodSpec (
149
+ containers = [
150
+ k8s .V1Container (
151
+ name = "base" ,
152
+ resources = k8s .V1ResourceRequirements (
153
+ requests = {
154
+ "cpu" : "100m" ,
155
+ "memory" : "384Mi" ,
156
+ },
157
+ limits = {
158
+ "cpu" : 1 ,
159
+ "memory" : "500Mi" ,
160
+ },
161
+ ),
162
+ )
163
+ ]
164
+ )
165
+ )
166
+ },
167
+ id = "pod_override_executor_config" ,
168
+ ),
169
+ ],
170
+ )
171
+ def test_celery_integration (self , broker_url , executor_config ):
141
172
from airflow .providers .celery .executors import celery_executor , celery_executor_utils
142
173
143
174
def fake_execute_workload (command ):
@@ -157,6 +188,7 @@ def fake_execute_workload(command):
157
188
try_number = 0 ,
158
189
priority_weight = 1 ,
159
190
queue = celery_executor_utils .celery_configuration ["task_default_queue" ],
191
+ executor_config = executor_config ,
160
192
)
161
193
keys = [
162
194
TaskInstanceKey ("id" , "success" , "abc" , 0 , - 1 ),
0 commit comments