diff --git a/WifiConnection.py b/NetworkPrinterOutputDevice.py
similarity index 68%
rename from WifiConnection.py
rename to NetworkPrinterOutputDevice.py
index 996cb2e5ea..de5a5961fd 100644
--- a/WifiConnection.py
+++ b/NetworkPrinterOutputDevice.py
@@ -1,19 +1,18 @@
-from UM.OutputDevice.OutputDevice import OutputDevice
-from UM.OutputDevice import OutputDeviceError
 import threading
-import json
 import time
-import base64
 import requests
 
-from . import HttpUploadDataStream
 from UM.i18n import i18nCatalog
-from UM.Signal import Signal, SignalEmitter
+from UM.Signal import Signal
 from UM.Application import Application
 from UM.Logger import Logger
+
+from cura.PrinterOutputDevice import PrinterOutputDevice, ConnectionState
+
 i18n_catalog = i18nCatalog("cura")
 
-class WifiConnection(OutputDevice, SignalEmitter):
+
+class NetworkPrinterOutputDevice(PrinterOutputDevice):
     def __init__(self, key, address, info):
         super().__init__(key)
         self._address = address
@@ -27,8 +26,6 @@ class WifiConnection(OutputDevice, SignalEmitter):
 
         self._json_printer_state = None
 
-        self._is_connected = False
-
         self._api_version = "1"
         self._api_prefix = "/api/v" + self._api_version + "/"
         self.setName(key)
@@ -36,38 +33,20 @@ class WifiConnection(OutputDevice, SignalEmitter):
         self.setDescription(i18n_catalog.i18nc("@info:tooltip", "Print with WIFI"))
         self.setIconName("print")
 
-    connectionStateChanged = Signal()
-
     def getKey(self):
         return self._key
 
-    def isConnected(self):
-        return self._is_connected
-
-    ##  Set the connection state of this connection.
-    #   Although we use a restfull API, we do poll the api to check if the machine is still responding.
-    def setConnectionState(self, state):
-        if state != self._is_connected:
-            Logger.log("i", "setting connection state of %s to %s " %(self._address, state))
-            self._is_connected = state
-            self.connectionStateChanged.emit(self._address)
-        else:
-            self._is_connected = state
-
     def _update(self):
-        while self._thread:
-            self.setConnectionState(True)
+        while self._connection_state == ConnectionState.connected or self._connection_state == ConnectionState.busy:
             try:
                 reply = self._httpGet("printer")
                 if reply.status_code == 200:
                     self._json_printer_state = reply.json()
-                    if not self._is_connected:
-                        self.setConnectionState(True)
                 else:
-                    self.setConnectionState(False)
+                    self.setConnectionState(ConnectionState.error)
             except:
-                self.setConnectionState(False)
-            time.sleep(1)
+                self.setConnectionState(ConnectionState.error)
+            time.sleep(1)  # Poll every second for printer state.
 
     def close(self):
         self._do_update = False
diff --git a/WifiOutputDevicePlugin.py b/WifiOutputDevicePlugin.py
index 7b89002a43..8d4ea516e8 100644
--- a/WifiOutputDevicePlugin.py
+++ b/WifiOutputDevicePlugin.py
@@ -1,5 +1,5 @@
 from UM.OutputDevice.OutputDevicePlugin import OutputDevicePlugin
-from . import WifiConnection
+from . import NetworkPrinterOutputDevice
 
 from zeroconf import Zeroconf, ServiceBrowser, ServiceStateChange
 from UM.Signal import Signal, SignalEmitter
@@ -35,7 +35,7 @@ class WifiOutputDevicePlugin(OutputDevicePlugin, SignalEmitter):
 
     ##  Because the model needs to be created in the same thread as the QMLEngine, we use a signal.
     def addConnection(self, name, address, properties):
-        connection = WifiConnection.WifiConnection(name, address, properties)
+        connection = NetworkPrinterOutputDevice.NetworkPrinterOutputDevice(name, address, properties)
         self._connections[address] = connection
         if connection.getKey() == Application.getInstance().getMachineManager().getActiveMachineInstance().getKey():
             self._connections[address].connect()
@@ -60,4 +60,5 @@ class WifiOutputDevicePlugin(OutputDevicePlugin, SignalEmitter):
 
         elif state_change == ServiceStateChange.Removed:
             info = zeroconf.get_service_info(service_type, name)
-            address = '.'.join(map(lambda n: str(n), info.address))
+            if info:
+                address = '.'.join(map(lambda n: str(n), info.address))