Skip to content

Commit f101567

Browse files
tieneupintieneupin
committed
Add integrated test for the FIB milling progress workflow
1 parent 6b57da5 commit f101567

1 file changed

Lines changed: 388 additions & 0 deletions

File tree

‎tests/workflows/fib/test_register_milling_progress.py‎

Lines changed: 388 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,388 @@
1+
from pathlib import Path
2+
from unittest.mock import MagicMock
3+
4+
from pytest_mock import MockerFixture
5+
from sqlalchemy.orm import Session as SQLAlchemySession
6+
from sqlmodel import Session as SQLModelSession, select as sm_select
7+
8+
import murfey.util.db as MurfeyDB
9+
from murfey.workflows.fib.register_milling_progress import run
10+
from tests.conftest import ExampleVisit
11+
12+
# Module-wide variables
13+
session_id = 10
14+
visit_name = f"{ExampleVisit.proposal_code}{ExampleVisit.proposal_number}-{ExampleVisit.visit_number}"
15+
instrument_name = ExampleVisit.instrument_name
16+
17+
18+
# Construct test FIB AutoTEM site info for reading (copied and adapted from actual session)
19+
site_info = {
20+
"project_name": visit_name,
21+
"site_name": "Lamella",
22+
"site_number": 1,
23+
"stage_info": {
24+
"preparation_site": {
25+
"x": -0.0031091224743875403,
26+
"y": 0.00420867925495798,
27+
"z": 0.0323644854106331,
28+
"rotation": 285.003247202109,
29+
"tilt_alpha": 25.9996646026832,
30+
"slot_number": 1,
31+
},
32+
"chunk_site": {
33+
"x": -0.0030037500000000003,
34+
"y": 0.004293,
35+
"z": 0.032350405092592606,
36+
"rotation": 285.003247202109,
37+
"tilt_alpha": -0.000134158926728586,
38+
"slot_number": 1,
39+
},
40+
"thinning_site": {
41+
"x": -0.0030037500000000003,
42+
"y": 0.004293,
43+
"z": 0.032350405092592606,
44+
"rotation": 285.003247202109,
45+
"tilt_alpha": -0.000134158926728586,
46+
"slot_number": 1,
47+
},
48+
"chunk_coincidence_params": {
49+
"x": -0.0030048260286678298,
50+
"y": 0.004308828126160981,
51+
"z": 0.0323400707790533,
52+
"rotation": 285.003247202109,
53+
"tilt_alpha": -0.000134158926728586,
54+
"slot_number": 1,
55+
},
56+
"thinning_params": {
57+
"x": -0.0030037500000000003,
58+
"y": 0.004293,
59+
"z": 0.032350405092592606,
60+
"rotation": 285.003247202109,
61+
"tilt_alpha": -0.000134158926728586,
62+
"slot_number": 1,
63+
},
64+
},
65+
"steps": {
66+
"eucentric_tilt": {
67+
"step_name": "Eucentric Tilt",
68+
"recipe_name": "Preparation",
69+
"is_enabled": False,
70+
"status": "Finished",
71+
"execution_time": 328.6657458,
72+
},
73+
"artificial_features": {
74+
"step_name": "Artificial Features",
75+
"recipe_name": "Preparation",
76+
"is_enabled": False,
77+
"execution_time": 0.0,
78+
"beam_type": "Ion",
79+
"voltage": 30000.0,
80+
"current": 5e-10,
81+
},
82+
"milling_angle": {
83+
"step_name": "Milling Angle",
84+
"recipe_name": "Preparation",
85+
"is_enabled": False,
86+
"status": "Finished",
87+
"execution_time": 208.4942922,
88+
"milling_angle": 12.0,
89+
},
90+
"image_acquisition": {
91+
"step_name": "Image Acquisition",
92+
"recipe_name": "Preparation",
93+
"is_enabled": False,
94+
"status": "Finished",
95+
"execution_time": 19.7872887,
96+
"site_location_type": "Chunk",
97+
},
98+
"lamella_placement": {
99+
"step_name": "Lamella Placement",
100+
"recipe_name": "Preparation",
101+
"is_enabled": False,
102+
"execution_time": 0.0,
103+
},
104+
"delay_1": {
105+
"step_name": "Delay",
106+
"recipe_name": "Milling",
107+
"is_enabled": False,
108+
"execution_time": 0.0,
109+
},
110+
"reference_definition": {
111+
"step_name": "Reference Definition",
112+
"recipe_name": "Milling",
113+
"is_enabled": False,
114+
"status": "Finished",
115+
"execution_time": 75.6320442,
116+
"site_location_type": "Chunk",
117+
},
118+
"reference_definition_electron": {
119+
"step_name": "Electron Reference Definition",
120+
"recipe_name": "Milling",
121+
"is_enabled": False,
122+
"status": "Finished",
123+
"execution_time": 131.7462301,
124+
},
125+
"stress_relief_cuts": {
126+
"step_name": "Stress Relief Cuts",
127+
"recipe_name": "Milling",
128+
"is_enabled": False,
129+
"execution_time": 0.0,
130+
"beam_type": "Ion",
131+
"voltage": 30000.0,
132+
"current": 1e-09,
133+
"depth_correction": 3.0,
134+
},
135+
"reference_redefinition_1": {
136+
"step_name": "Reference Redefinition 1",
137+
"recipe_name": "Milling",
138+
"is_enabled": False,
139+
"execution_time": 0.0,
140+
"site_location_type": "Chunk",
141+
},
142+
"rough_milling": {
143+
"step_name": "Rough Milling",
144+
"recipe_name": "Milling",
145+
"is_enabled": False,
146+
"status": "Finished",
147+
"execution_time": 1.5929074719,
148+
"beam_type": "Ion",
149+
"voltage": 30000.0,
150+
"current": 1e-09,
151+
"depth_correction": 3.0,
152+
"lamella_offset": 2e-06,
153+
"trench_height_front": 5.6338138462765e-06,
154+
"trench_height_rear": 8.60473362403628e-06,
155+
"width_overlap_front_left": 2e-06,
156+
"width_overlap_front_right": 2e-06,
157+
"width_overlap_rear_left": 2e-06,
158+
"width_overlap_rear_right": 2e-06,
159+
},
160+
"rough_milling_electron": {
161+
"step_name": "Rough Milling - Electron Image",
162+
"recipe_name": "Milling",
163+
"is_enabled": False,
164+
"execution_time": 0.0,
165+
"site_location_type": "Chunk",
166+
"beam_type": "Electron",
167+
"voltage": 2000.0,
168+
"current": 1.25e-11,
169+
},
170+
"reference_redefinition_2": {
171+
"step_name": "Reference Redefinition 2",
172+
"recipe_name": "Milling",
173+
"is_enabled": False,
174+
"status": "Finished",
175+
"execution_time": 105.1459769,
176+
"site_location_type": "Chunk",
177+
},
178+
"medium_milling": {
179+
"step_name": "Medium Milling",
180+
"recipe_name": "Milling",
181+
"is_enabled": False,
182+
"status": "Finished",
183+
"execution_time": 298.1471377,
184+
"beam_type": "Ion",
185+
"voltage": 30000.0,
186+
"current": 5e-10,
187+
"depth_correction": 2.0,
188+
"lamella_offset": 1.5e-06,
189+
"width_overlap_front_left": 1.5e-06,
190+
"width_overlap_front_right": 1.5e-06,
191+
"width_overlap_rear_left": 1.5e-06,
192+
"width_overlap_rear_right": 1.5e-06,
193+
},
194+
"medium_milling_electron": {
195+
"step_name": "Medium Milling - Electron Image",
196+
"recipe_name": "Milling",
197+
"is_enabled": False,
198+
"execution_time": 0.0,
199+
"site_location_type": "Chunk",
200+
"beam_type": "Electron",
201+
"voltage": 2000.0,
202+
"current": 1.25e-11,
203+
},
204+
"fine_milling": {
205+
"step_name": "Fine Milling",
206+
"recipe_name": "Milling",
207+
"is_enabled": False,
208+
"status": "Finished",
209+
"execution_time": 397.414388,
210+
"beam_type": "Ion",
211+
"voltage": 30000.0,
212+
"current": 3e-10,
213+
"depth_correction": 2.0,
214+
"lamella_offset": 1e-06,
215+
"width_overlap_front_left": 1e-06,
216+
"width_overlap_front_right": 1e-06,
217+
"width_overlap_rear_left": 1e-06,
218+
"width_overlap_rear_right": 1e-06,
219+
},
220+
"fine_milling_electron": {
221+
"step_name": "Fine Milling - Electron Image",
222+
"recipe_name": "Milling",
223+
"is_enabled": False,
224+
"execution_time": 0.0,
225+
"site_location_type": "Chunk",
226+
"beam_type": "Electron",
227+
"voltage": 2000.0,
228+
"current": 1.25e-11,
229+
},
230+
"finer_milling": {
231+
"step_name": "Finer Milling",
232+
"recipe_name": "Milling",
233+
"is_enabled": False,
234+
"status": "Finished",
235+
"execution_time": 887.9686893,
236+
"beam_type": "Ion",
237+
"voltage": 30000.0,
238+
"current": 1e-10,
239+
"depth_correction": 2.0,
240+
"lamella_offset": 4.0000000000000003e-07,
241+
"width_overlap_front_left": 5.000000000000001e-07,
242+
"width_overlap_front_right": 5.000000000000001e-07,
243+
"width_overlap_rear_left": 5.000000000000001e-07,
244+
"width_overlap_rear_right": 5.000000000000001e-07,
245+
},
246+
"finer_milling_electron": {
247+
"step_name": "Finer Milling - Electron Image",
248+
"recipe_name": "Milling",
249+
"is_enabled": False,
250+
"status": "Finished",
251+
"execution_time": 59.2296577,
252+
"site_location_type": "Chunk",
253+
"beam_type": "Electron",
254+
"voltage": 2000.0,
255+
"current": 2.5e-11,
256+
},
257+
"delay_2": {
258+
"step_name": "Delay",
259+
"recipe_name": "Thinning",
260+
"is_enabled": False,
261+
"execution_time": 0.0,
262+
},
263+
"polishing_1": {
264+
"step_name": "Polishing 1",
265+
"recipe_name": "Thinning",
266+
"is_enabled": False,
267+
"status": "Finished",
268+
"execution_time": 680.5567315,
269+
"beam_type": "Ion",
270+
"voltage": 30000.0,
271+
"current": 5e-11,
272+
"depth_correction": 2.0,
273+
"lamella_offset": 2.5000000000000004e-07,
274+
"width_overlap_front_left": 0.0,
275+
"width_overlap_front_right": 0.0,
276+
"width_overlap_rear_left": 0.0,
277+
"width_overlap_rear_right": 0.0,
278+
},
279+
"polishing_1_electron": {
280+
"step_name": "Polishing 1 - Electron Image",
281+
"recipe_name": "Thinning",
282+
"is_enabled": False,
283+
"execution_time": 0.0,
284+
"site_location_type": "Thinning",
285+
"beam_type": "Electron",
286+
"voltage": 2000.0,
287+
"current": 1.25e-11,
288+
},
289+
"polishing_2": {
290+
"step_name": "Polishing 2",
291+
"recipe_name": "Thinning",
292+
"is_enabled": False,
293+
"status": "Finished",
294+
"execution_time": 1.170488927,
295+
"beam_type": "Ion",
296+
"voltage": 30000.0,
297+
"current": 3e-11,
298+
"depth_correction": 1.5,
299+
"lamella_offset": 0.0,
300+
"width_overlap_front_left": 0.0,
301+
"width_overlap_front_right": 0.0,
302+
"width_overlap_rear_left": 0.0,
303+
"width_overlap_rear_right": 0.0,
304+
},
305+
"polishing_2_ion": {
306+
"step_name": "Polishing 2 - Ion Image",
307+
"recipe_name": "Thinning",
308+
"is_enabled": False,
309+
"execution_time": 0.0,
310+
"site_location_type": "Thinning",
311+
"beam_type": "Ion",
312+
"voltage": 30000.0,
313+
"current": 3e-11,
314+
},
315+
"polishing_2_electron": {
316+
"step_name": "Polishing 2 - Electron Image",
317+
"recipe_name": "Thinning",
318+
"is_enabled": False,
319+
"status": "Finished",
320+
"execution_time": 56.9180832,
321+
"site_location_type": "Thinning",
322+
"beam_type": "Electron",
323+
"voltage": 2000.0,
324+
"current": 2.5e-11,
325+
},
326+
},
327+
}
328+
# Construct the RabbitMQ message received
329+
message = {
330+
"register": "fib.register_milling_progress",
331+
"session_id": session_id,
332+
"site_info": site_info,
333+
}
334+
335+
336+
def test_run_with_db(
337+
mocker: MockerFixture,
338+
murfey_db_session: SQLModelSession,
339+
ispyb_db_session: SQLAlchemySession,
340+
mock_ispyb_credentials: Path,
341+
):
342+
# Add a test visit to the database
343+
if not (
344+
session_entry := murfey_db_session.exec(
345+
sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id)
346+
).one_or_none()
347+
):
348+
session_entry = MurfeyDB.Session(id=session_id)
349+
session_entry.name = visit_name
350+
session_entry.visit = visit_name
351+
session_entry.instrument_name = instrument_name
352+
353+
murfey_db_session.add(session_entry)
354+
murfey_db_session.commit()
355+
356+
# Mock the ISPyB connection where the TransportManager class is located
357+
mock_security_config = MagicMock()
358+
mock_security_config.ispyb_credentials = mock_ispyb_credentials
359+
mocker.patch(
360+
"murfey.server.ispyb.get_security_config",
361+
return_value=mock_security_config,
362+
)
363+
mocker.patch(
364+
"murfey.server.ispyb.ISPyBSession",
365+
return_value=ispyb_db_session,
366+
)
367+
368+
# Mock the ISPYB connection when registering DataCollectionGroup
369+
mocker.patch(
370+
"murfey.workflows.register_data_collection_group.ISPyBSession",
371+
return_value=ispyb_db_session,
372+
)
373+
374+
# Patch the TransportManager object in the workflows called
375+
from murfey.server.ispyb import TransportManager
376+
377+
mocker.patch(
378+
"murfey.workflows.register_data_collection_group._transport_object",
379+
new=TransportManager("PikaTransport"),
380+
)
381+
382+
# Run the workflow twice (fresh insert and update existing)
383+
for i in range(2):
384+
result = run(
385+
message=message,
386+
murfey_db=murfey_db_session,
387+
)
388+
assert result.get("success", False)

0 commit comments

Comments
 (0)