process_replay: refactor migrate_deviceState (#34864)

* Refactor migrate_deviceState

* Fix again

selfdrive/test/process_replay/migration.py

@@ -241,14 +241,16 @@ def migrate_gpsLocation(msgs):
 
 @migration(inputs=["deviceState", "initData"])
 def migrate_deviceState(msgs):
+  init_data = next((m.initData for _, m in msgs if m.which() == 'initData'), None)
+  device_state = next((m.deviceState for _, m in msgs if m.which() == 'deviceState'), None)
+  if init_data is None or device_state is None:
+    return [], [], []
+
   ops = []
-  dt = None
   for i, msg in msgs:
-    if msg.which() == 'initData':
-      dt = msg.initData.deviceType
     if msg.which() == 'deviceState':
       n = msg.as_builder()
-      n.deviceState.deviceType = dt
+      n.deviceState.deviceType = init_data.deviceType
       ops.append((i, n.as_reader()))
   return ops, [], []