@@ -155,7 +155,10 @@ async def jobs_search(
155155 await create_location (
156156 async_session ,
157157 location_kwargs = kwargs ,
158- address_kwargs = {"urls" : [random_string (32 ), random_string (32 )]}, # Each location has 2 addresses
158+ # Each location has 2 addresses
159+ address_kwargs = {
160+ "urls" : [random_string (32 ), random_string (32 )],
161+ },
159162 )
160163 for kwargs in location_kwargs
161164 ]
@@ -170,7 +173,9 @@ async def jobs_search(
170173 await create_location (
171174 async_session ,
172175 location_kwargs = {"type" : "random" },
173- address_kwargs = {"urls" : [random_string (32 ), random_string (32 )]},
176+ address_kwargs = {
177+ "urls" : [random_string (32 ), random_string (32 )],
178+ },
174179 )
175180 for _ in range (3 )
176181 ]
@@ -207,11 +212,61 @@ async def jobs_search(
207212
208213 jobs_by_name = {job .name : job for job in jobs_with_names }
209214 jobs_by_location = dict (
210- zip ([location .name for location in locations_with_names ], jobs_with_location_names , strict = False ),
215+ zip (
216+ [location .name for location in locations_with_names ],
217+ jobs_with_location_names ,
218+ strict = False ,
219+ ),
220+ )
221+ jobs_by_address = dict (
222+ zip (
223+ addresses_url ,
224+ [job for job in jobs_with_address_urls for _ in range (2 )],
225+ strict = False ,
226+ ),
211227 )
212- jobs_by_address = dict (zip (addresses_url , [job for job in jobs_with_address_urls for _ in range (2 )], strict = False ))
213228
214229 yield jobs_by_name , jobs_by_location , jobs_by_address
215230
216231 async with async_session_maker () as async_session :
217232 await clean_db (async_session )
233+
234+
235+ @pytest_asyncio .fixture
236+ async def jobs_with_locations_and_types (
237+ async_session_maker : Callable [[], AbstractAsyncContextManager [AsyncSession ]],
238+ ) -> AsyncGenerator [tuple [Job ]]:
239+ async with async_session_maker () as async_session :
240+ cluster_location = await create_location (async_session , location_kwargs = {"name" : "my-cluster" , "type" : "yarn" })
241+ airflow_location = await create_location (
242+ async_session ,
243+ location_kwargs = {"name" : "airflow-host" , "type" : "http" },
244+ )
245+ cluster_type = await create_job_type (async_session , {"type" : "SPARK_APPLICATION" })
246+ airflow_dag_type = await create_job_type (async_session , {"type" : "AIRFLOW_DAG" })
247+ airflow_task_type = await create_job_type (async_session , {"type" : "AIRFLOW_TASK" })
248+ cluster_job = await create_job (
249+ async_session ,
250+ location_id = cluster_location .id ,
251+ job_type_id = cluster_type .id ,
252+ job_kwargs = {"name" : "my-job_cluster" },
253+ )
254+ dag_job = await create_job (
255+ async_session ,
256+ location_id = airflow_location .id ,
257+ job_type_id = airflow_dag_type .id ,
258+ job_kwargs = {"name" : "my-job_dag" },
259+ )
260+ task_job = await create_job (
261+ async_session ,
262+ location_id = airflow_location .id ,
263+ job_type_id = airflow_task_type .id ,
264+ job_kwargs = {"name" : "my-job_task" },
265+ )
266+
267+ async_session .expunge_all ()
268+
269+ yield (cluster_job , dag_job , task_job )
270+
271+ async with async_session_maker () as async_session :
272+ await clean_db (async_session )
0 commit comments