Fix dict get values

sonic-net · Nov 13, 2024 · 1ab0206 · 1ab0206
1 parent a4340d2
commit 1ab0206
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/sonic-pcied/tests/test_DaemonPcied.py b/sonic-pcied/tests/test_DaemonPcied.py
@@ -149,13 +149,10 @@ def test_is_dpu_in_detaching_mode(self):
         daemon_pcied.detach_info = mock.MagicMock()
         daemon_pcied.detach_info.getKeys = mock.MagicMock(return_value=['DPU_0', 'DPU_1'])
         daemon_pcied.detach_info.get = mock.MagicMock(
-            side_effect=lambda key: detach_info.get(
-            key,
-            {'bus_info': '0000:03:00.1', 'dpu_state': 'detaching'}
-            ) if key == 'DPU_0' else detach_info.get(
-            key,
-            {'bus_info': '0000:03:00.2', 'dpu_state': 'attached'}
-            ) if key == 'DPU_1' else None
+            side_effect=lambda key: {
+                'DPU_0': {'bus_info': '0000:03:00.1', 'dpu_state': 'detaching'},
+                'DPU_1': {'bus_info': '0000:03:00.2', 'dpu_state': 'attached'}
+            }.get(key, None)
         )
 
         # Test when the device is in detaching mode