diff --git a/bec_widgets/line_plot.py b/bec_widgets/line_plot.py
index 71e2976f..33b310a0 100644
--- a/bec_widgets/line_plot.py
+++ b/bec_widgets/line_plot.py
@@ -1,14 +1,13 @@
-from typing import Any
+import os
 import warnings
+from typing import Any
 
 import numpy as np
 import pyqtgraph as pg
-from pyqtgraph import mkPen, mkBrush, mkColor
-
+from bec_lib import BECClient
+from pyqtgraph import mkBrush, mkColor, mkPen
 from pyqtgraph.Qt import QtCore, QtWidgets, uic
 from pyqtgraph.Qt.QtCore import pyqtSignal
-from bec_lib import BECClient
-import os
 
 
 class BasicPlot(QtWidgets.QWidget):
@@ -169,7 +168,7 @@ class BasicPlot(QtWidgets.QWidget):
         remove_y_value_index = [
             index
             for index, y_value in enumerate(self.y_value_list)
-            if y_value not in client.device_manager.devices.keys()
+            if y_value not in client.device_manager.devices
         ]
         if remove_y_value_index:
             for ii in sorted(remove_y_value_index, reverse=True):
@@ -252,6 +251,7 @@ class BasicPlot(QtWidgets.QWidget):
 
 if __name__ == "__main__":
     import argparse
+
     from bec_widgets import ctrl_c
 
     parser = argparse.ArgumentParser()
diff --git a/tests/test_line_plot.py b/tests/test_line_plot.py
index 53aebb5e..41506dac 100644
--- a/tests/test_line_plot.py
+++ b/tests/test_line_plot.py
@@ -1,7 +1,9 @@
-from pytestqt import qtbot
 from unittest import mock
-from bec_widgets import line_plot
+
 import numpy as np
+from pytestqt import qtbot
+
+from bec_widgets import line_plot
 
 
 def test_line_plot_emits_no_signal(qtbot):
@@ -63,12 +65,14 @@ def test_line_plot_raise_warning_wrong_signal_request(qtbot):
     metadata = {"scanID": "test", "scan_number": 1, "scan_report_devices": ["x"]}
     with mock.patch("bec_widgets.line_plot.BECClient") as mock_client:
         # TODO fix mock_client
-        mock_client.device_manager.devices.keys.return_value = ["y1"]
+        mock_dict = {"y1": [1, 2]}
+        mock_client().device_manager.devices.__contains__.side_effect = mock_dict.__contains__
+
         # = {"y1": [1, 2]}
         with mock.patch.object(plot, "update_signal") as mock_update_signal:
             mock_update_signal.emit()
             plot(data=data, metadata=metadata)
-            assert all(plot.y_value_list == ["y1"])
+            assert plot.y_value_list == ["y1"]
 
 
 def test_line_plot_update(qtbot):