diff --git a/custom_components/nvidia_shield/config_flow.py b/custom_components/nvidia_shield/config_flow.py
index 91c479d..f9d7c1c 100644
--- a/custom_components/nvidia_shield/config_flow.py
+++ b/custom_components/nvidia_shield/config_flow.py
@@ -8,7 +8,8 @@ from homeassistant.core import callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers import selector
 
-from .adb_session import AdbMediaSessionPoller, ensure_adb_key
+from .adb_session import ensure_adb_key
+from .connection import ShieldConnectionManager
 from .const import (
     ATTR_MEDIA_DAG_INPUT_MAP,
     CONF_ADB_ENABLED,
@@ -207,10 +208,9 @@ class NvidiaShieldConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     ensure_adb_key, self.hass.config.config_dir
                 )
 
-                poller = AdbMediaSessionPoller(host, port, key_path)
-                connected = await poller.connect()
-                await poller.disconnect()
-
+                connected = await ShieldConnectionManager.test_connect(
+                    self.hass, host, port, key_path
+                )
                 if connected:
                     return self._create_entry()
                 else: