diff --git a/hyperglass/execution/drivers/__init__.py b/hyperglass/execution/drivers/__init__.py index 8abe2b7..fd3bb6f 100644 --- a/hyperglass/execution/drivers/__init__.py +++ b/hyperglass/execution/drivers/__init__.py @@ -1,6 +1,5 @@ """Individual transport driver classes & subclasses.""" -# Project -from hyperglass.execution.drivers.agent import AgentConnection -from hyperglass.execution.drivers.ssh_netmiko import NetmikoConnection -from hyperglass.execution.drivers.ssh_scrapli import ScrapliConnection +from .agent import AgentConnection +from .ssh_netmiko import NetmikoConnection +from .ssh_scrapli import ScrapliConnection diff --git a/hyperglass/execution/construct.py b/hyperglass/execution/drivers/_construct.py similarity index 98% rename from hyperglass/execution/construct.py rename to hyperglass/execution/drivers/_construct.py index aaccfca..81015ed 100644 --- a/hyperglass/execution/construct.py +++ b/hyperglass/execution/drivers/_construct.py @@ -12,7 +12,7 @@ from operator import attrgetter # Project from hyperglass.log import log -from hyperglass.constants import ( +from hyperglass.constants import ( # APPEND_NEWLINE, TRANSPORT_REST, TARGET_FORMAT_SPACE, TARGET_JUNIPER_ASPATH, diff --git a/hyperglass/execution/drivers/agent.py b/hyperglass/execution/drivers/agent.py index a72b7b2..d1d97d8 100644 --- a/hyperglass/execution/drivers/agent.py +++ b/hyperglass/execution/drivers/agent.py @@ -19,7 +19,8 @@ from hyperglass.util import parse_exception from hyperglass.encode import jwt_decode, jwt_encode from hyperglass.exceptions import RestError, ResponseEmpty from hyperglass.configuration import params -from hyperglass.execution.drivers._common import Connection + +from ._common import Connection class AgentConnection(Connection): diff --git a/hyperglass/execution/drivers/ssh.py b/hyperglass/execution/drivers/ssh.py index 9e226ae..ce9713b 100644 --- a/hyperglass/execution/drivers/ssh.py +++ b/hyperglass/execution/drivers/ssh.py @@ -8,7 +8,8 @@ from hyperglass.log import log from hyperglass.exceptions import ScrapeError from hyperglass.configuration import params from hyperglass.compat._sshtunnel import BaseSSHTunnelForwarderError, open_tunnel -from hyperglass.execution.drivers._common import Connection + +from ._common import Connection class SSHConnection(Connection): diff --git a/hyperglass/execution/drivers/ssh_scrapli.py b/hyperglass/execution/drivers/ssh_scrapli.py index 69ab771..7b21108 100644 --- a/hyperglass/execution/drivers/ssh_scrapli.py +++ b/hyperglass/execution/drivers/ssh_scrapli.py @@ -32,7 +32,8 @@ from hyperglass.exceptions import ( UnsupportedDevice, ) from hyperglass.configuration import params -from hyperglass.execution.drivers.ssh import SSHConnection + +from .ssh import SSHConnection SCRAPLI_DRIVER_MAP = { "cisco_ios": AsyncIOSXEDriver, @@ -90,7 +91,7 @@ class ScrapliConnection(SSHConnection): responses = () async with driver as connection: - + await connection.get_prompt() for query in self.query: raw = await connection.send_command(query) responses += (raw.result,) diff --git a/hyperglass/execution/main.py b/hyperglass/execution/main.py index 046c221..4d9cb6a 100644 --- a/hyperglass/execution/main.py +++ b/hyperglass/execution/main.py @@ -14,13 +14,10 @@ from typing import Any, Dict, Union, Callable from hyperglass.log import log from hyperglass.util import validate_nos from hyperglass.exceptions import DeviceTimeout, ResponseEmpty +from hyperglass.models.api import Query from hyperglass.configuration import params -from hyperglass.api.models.query import Query -from hyperglass.execution.drivers import ( - AgentConnection, - NetmikoConnection, - ScrapliConnection, -) + +from .drivers import AgentConnection, NetmikoConnection, ScrapliConnection DRIVER_MAP = { "scrapli": ScrapliConnection,