Fix DiscoveredMachineModel test

cura/UI/MachineModels/DiscoveredPrintersModel.py

@@ -11,6 +11,8 @@ from UM.Logger import Logger
 if TYPE_CHECKING:
     from PyQt5.QtCore import QObject
 
+    from cura.PrinterOutput.NetworkedPrinterOutputDevice import NetworkedPrinterOutputDevice
+
 
 catalog = i18nCatalog("cura")
 
@@ -18,7 +20,7 @@ catalog = i18nCatalog("cura")
 class DiscoveredPrinter(QObject):
 
     def __init__(self, ip_address: str, key: str, name: str, create_callback: Callable[[str], None], machine_type: str,
-                 device, parent = None) -> None:
+                 device: "NetworkedPrinterOutputDevice", parent: Optional["QObject"] = None) -> None:
         super().__init__(parent)
 
         self._ip_address = ip_address
@@ -67,7 +69,7 @@ class DiscoveredPrinter(QObject):
         return self.readable_machine_type.lower() == "unknown"
 
     @pyqtProperty(QObject, constant = True)
-    def device(self):
+    def device(self) -> "NetworkedPrinterOutputDevice":
         return self._device
 
 
@@ -78,7 +80,7 @@ class DiscoveredPrinter(QObject):
 #
 class DiscoveredPrintersModel(QObject):
 
-    def __init__(self, parent: Optional["QObject"]) -> None:
+    def __init__(self, parent: Optional["QObject"] = None) -> None:
         super().__init__(parent)
 
         self._discovered_printer_by_ip_dict = dict()  # type: Dict[str, DiscoveredPrinter]
@@ -92,7 +94,7 @@ class DiscoveredPrintersModel(QObject):
         return item_list
 
     def addDiscoveredPrinter(self, ip_address: str, key: str, name: str, create_callback: Callable[[str], None],
-                             machine_type: str, device) -> None:
+                             machine_type: str, device: "NetworkedPrinterOutputDevice") -> None:
         if ip_address in self._discovered_printer_by_ip_dict:
             Logger.log("e", "Printer with ip [%s] has already been added", ip_address)
             return

tests/TestMachineManager.py

@@ -45,20 +45,3 @@ def test_setActiveMachine(machine_manager):
 
             # Although we mocked the application away, we still want to know if it was notified about the attempted change.
             machine_manager._application.setGlobalContainerStack.assert_called_with(mocked_global_stack)
-
-
-def test_discoveredMachine(machine_manager):
-    mocked_callback = MagicMock()
-    machine_manager.addDiscoveredPrinter("test", "zomg", mocked_callback, "derp")
-    machine_manager.addMachineFromDiscoveredPrinter("test")
-    mocked_callback.assert_called_with("test")
-
-    assert len(machine_manager.discoveredPrinters) == 1
-
-    # Test if removing it works
-    machine_manager.removeDiscoveredPrinter("test")
-    assert len(machine_manager.discoveredPrinters) == 0
-
-    # Just in case, nothing should happen.
-    machine_manager.addMachineFromDiscoveredPrinter("test")
-    assert mocked_callback.call_count == 1

tests/UI/MachineModels/TestDiscoveredPrintersModel.py

@@ -0,0 +1,27 @@
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from cura.UI.MachineModels.DiscoveredPrintersModel import DiscoveredPrintersModel
+
+
+@pytest.fixture()
+def discovered_printer_model(application) -> DiscoveredPrintersModel:
+    return DiscoveredPrintersModel()
+
+
+def test_discoveredPrinters(discovered_printer_model):
+    mocked_device = MagicMock()
+
+    mocked_callback = MagicMock()
+    discovered_printer_model.addDiscoveredPrinter("ip", "key", "name", mocked_callback, "machine_type", mocked_device)
+    device = discovered_printer_model.discovered_printers[0]
+    discovered_printer_model.createMachineFromDiscoveredPrinter(device)
+    mocked_callback.assert_called_with("key")
+
+    assert len(discovered_printer_model.discovered_printers) == 1
+
+    # Test if removing it works
+    discovered_printer_model.removeDiscoveredPrinter("ip")
+    assert len(discovered_printer_model.discovered_printers) == 0