Fix network setup for collections in ansible-test.

7 files changed, 104 insertions, 10 deletions

diff --git a/test/lib/ansible_test/_data/completion/network.txt b/test/lib/ansible_test/_data/completion/network.txt
index 502388338a..cd762feb2a 100644
--- a/test/lib/ansible_test/_data/completion/network.txt
+++ b/test/lib/ansible_test/_data/completion/network.txt
@@ -1,4 +1,2 @@
-ios/csr1000v
-junos/vmx
-junos/vsrx
-vyos/1.1.8
+ios/csr1000v connection=local
+vyos/1.1.8 connection=local
diff --git a/test/lib/ansible_test/_internal/cli.py b/test/lib/ansible_test/_internal/cli.py
index 9e86ccb9d8..4ada06c043 100644
--- a/test/lib/ansible_test/_internal/cli.py
+++ b/test/lib/ansible_test/_internal/cli.py
@@ -19,6 +19,7 @@ from .util import (
     display,
     raw_command,
     get_docker_completion,
+    get_network_completion,
     get_remote_completion,
     generate_pip_command,
     read_lines_without_comments,
@@ -209,6 +210,10 @@ def parse_args():
     else:
         epilog = 'Install the "argcomplete" python package to enable tab completion.'
 
+    def key_value_type(value):  # type: (str) -> t.Tuple[str, str]
+        """Wrapper around key_value."""
+        return key_value(argparse, value)
+
     parser = argparse.ArgumentParser(epilog=epilog)
 
     common = argparse.ArgumentParser(add_help=False)
@@ -417,6 +422,18 @@ def parse_args():
                                      action='append',
                                      help='network platform/version').completer = complete_network_platform
 
+    network_integration.add_argument('--platform-collection',
+                                     type=key_value_type,
+                                     metavar='PLATFORM=COLLECTION',
+                                     action='append',
+                                     help='collection used to test platform').completer = complete_network_platform_collection
+
+    network_integration.add_argument('--platform-connection',
+                                     type=key_value_type,
+                                     metavar='PLATFORM=CONNECTION',
+                                     action='append',
+                                     help='connection used to test platform').completer = complete_network_platform_connection
+
     network_integration.add_argument('--inventory',
                                      metavar='PATH',
                                      help='path to inventory used for tests')
@@ -639,6 +656,16 @@ def parse_args():
     return args
 
 
+def key_value(argparse, value):  # type: (argparse_module, str) -> t.Tuple[str, str]
+    """Type parsing and validation for argparse key/value pairs separated by an '=' character."""
+    parts = value.split('=')
+
+    if len(parts) != 2:
+        raise argparse.ArgumentTypeError('"%s" must be in the format "key=value"' % value)
+
+    return tuple(parts)
+
+
 # noinspection PyProtectedMember
 def add_coverage_analyze(coverage_subparsers, coverage_common):  # type: (argparse_module._SubParsersAction, argparse_module.ArgumentParser) -> None
     """Add the `coverage analyze` subcommand."""
@@ -1039,11 +1066,35 @@ def complete_network_platform(prefix, parsed_args, **_):
     :type parsed_args: any
     :rtype: list[str]
     """
-    images = read_lines_without_comments(os.path.join(ANSIBLE_TEST_DATA_ROOT, 'completion', 'network.txt'), remove_blank_lines=True)
+    images = sorted(get_network_completion())
 
     return [i for i in images if i.startswith(prefix) and (not parsed_args.platform or i not in parsed_args.platform)]
 
 
+def complete_network_platform_collection(prefix, parsed_args, **_):
+    """
+    :type prefix: unicode
+    :type parsed_args: any
+    :rtype: list[str]
+    """
+    left = prefix.split('=')[0]
+    images = sorted(set(image.split('/')[0] for image in get_network_completion()))
+
+    return [i + '=' for i in images if i.startswith(left) and (not parsed_args.platform_collection or i not in [x[0] for x in parsed_args.platform_collection])]
+
+
+def complete_network_platform_connection(prefix, parsed_args, **_):
+    """
+    :type prefix: unicode
+    :type parsed_args: any
+    :rtype: list[str]
+    """
+    left = prefix.split('=')[0]
+    images = sorted(set(image.split('/')[0] for image in get_network_completion()))
+
+    return [i + '=' for i in images if i.startswith(left) and (not parsed_args.platform_connection or i not in [x[0] for x in parsed_args.platform_connection])]
+
+
 def complete_network_testcase(prefix, parsed_args, **_):
     """
     :type prefix: unicode
diff --git a/test/lib/ansible_test/_internal/config.py b/test/lib/ansible_test/_internal/config.py
index 293a6cf7ef..3431264f67 100644
--- a/test/lib/ansible_test/_internal/config.py
+++ b/test/lib/ansible_test/_internal/config.py
@@ -313,6 +313,8 @@ class NetworkIntegrationConfig(IntegrationConfig):
         super(NetworkIntegrationConfig, self).__init__(args, 'network-integration')
 
         self.platform = args.platform  # type: t.List[str]
+        self.platform_collection = dict(args.platform_collection or [])  # type: t.Dict[str, str]
+        self.platform_connection = dict(args.platform_connection or [])  # type: t.Dict[str, str]
         self.inventory = args.inventory  # type: str
         self.testcase = args.testcase  # type: str
 
diff --git a/test/lib/ansible_test/_internal/executor.py b/test/lib/ansible_test/_internal/executor.py
index b1e18f328f..3d1a1ab3c0 100644
--- a/test/lib/ansible_test/_internal/executor.py
+++ b/test/lib/ansible_test/_internal/executor.py
@@ -62,6 +62,7 @@ from .util import (
     generate_pip_command,
     find_python,
     get_docker_completion,
+    get_network_settings,
     get_remote_completion,
     cmd_quote,
     ANSIBLE_LIB_ROOT,
@@ -560,10 +561,12 @@ def network_inventory(remotes):
             ansible_host=remote.connection.hostname,
             ansible_user=remote.connection.username,
             ansible_ssh_private_key_file=os.path.abspath(remote.ssh_key.key),
-            ansible_network_os=remote.platform,
-            ansible_connection='local'
         )
 
+        settings = get_network_settings(remote.args, remote.platform, remote.version)
+
+        options.update(settings.inventory_vars)
+
         groups[remote.platform].append(
             '%s %s' % (
                 remote.name.replace('.', '-'),
@@ -1401,7 +1404,7 @@ def command_integration_role(args, target, start_at_task, test_dir, inventory_pa
             win_output_dir=r'C:\ansible_testing',
         ))
     elif isinstance(args, NetworkIntegrationConfig):
-        hosts = target.name[:target.name.find('_')]
+        hosts = target.network_platform
         gather_facts = False
     else:
         hosts = 'testhost'
diff --git a/test/lib/ansible_test/_internal/manage_ci.py b/test/lib/ansible_test/_internal/manage_ci.py
index 3c7446bafb..8cac2bef66 100644
--- a/test/lib/ansible_test/_internal/manage_ci.py
+++ b/test/lib/ansible_test/_internal/manage_ci.py
@@ -12,6 +12,7 @@ from .util import (
     cmd_quote,
     display,
     ANSIBLE_TEST_DATA_ROOT,
+    get_network_settings,
 )
 
 from .util_common import (
@@ -149,11 +150,14 @@ class ManageNetworkCI:
 
     def wait(self):
         """Wait for instance to respond to ansible ping."""
+        settings = get_network_settings(self.core_ci.args, self.core_ci.platform, self.core_ci.version)
+
         extra_vars = [
             'ansible_host=%s' % self.core_ci.connection.hostname,
             'ansible_port=%s' % self.core_ci.connection.port,
-            'ansible_connection=local',
             'ansible_ssh_private_key_file=%s' % self.core_ci.ssh_key.key,
+        ] + [
+            '%s=%s' % (key, value) for key, value in settings.inventory_vars.items()
         ]
 
         name = '%s-%s' % (self.core_ci.platform, self.core_ci.version.replace('.', '-'))
@@ -161,7 +165,7 @@ class ManageNetworkCI:
         env = ansible_environment(self.core_ci.args)
         cmd = [
             'ansible',
-            '-m', '%s_command' % self.core_ci.platform,
+            '-m', '%s%s_command' % (settings.collection + '.' if settings.collection else '', self.core_ci.platform),
             '-a', 'commands=?',
             '-u', self.core_ci.connection.username,
             '-i', '%s,' % name,
diff --git a/test/lib/ansible_test/_internal/target.py b/test/lib/ansible_test/_internal/target.py
index fa7526cb63..c3ba81268d 100644
--- a/test/lib/ansible_test/_internal/target.py
+++ b/test/lib/ansible_test/_internal/target.py
@@ -643,6 +643,10 @@ class IntegrationTarget(CompletionTarget):
         self.needs_file = tuple(sorted(set('/'.join(g.split('/')[2:]) for g in groups if
                                            g.startswith('needs/file/') and not g.startswith('needs/file/%s/' % targets_relative_path))))
 
+        # network platform
+        networks = [g.split('/')[1] for g in groups if g.startswith('network/')]
+        self.network_platform = networks[0] if networks else None
+
         for group in itertools.islice(groups, 0, len(groups)):
             if '/' in group:
                 parts = group.split('/')
diff --git a/test/lib/ansible_test/_internal/util.py b/test/lib/ansible_test/_internal/util.py
index 2b0218fb38..4a6ddfa1da 100644
--- a/test/lib/ansible_test/_internal/util.py
+++ b/test/lib/ansible_test/_internal/util.py
@@ -64,6 +64,7 @@ except AttributeError:
 
 DOCKER_COMPLETION = {}  # type: t.Dict[str, t.Dict[str, str]]
 REMOTE_COMPLETION = {}  # type: t.Dict[str, t.Dict[str, str]]
+NETWORK_COMPLETION = {}  # type: t.Dict[str, t.Dict[str, str]]
 PYTHON_PATHS = {}  # type: t.Dict[str, str]
 
 try:
@@ -134,6 +135,13 @@ def get_remote_completion():
     return get_parameterized_completion(REMOTE_COMPLETION, 'remote')
 
 
+def get_network_completion():
+    """
+    :rtype: dict[str, dict[str, str]]
+    """
+    return get_parameterized_completion(NETWORK_COMPLETION, 'network')
+
+
 def get_parameterized_completion(cache, name):
     """
     :type cache: dict[str, dict[str, str]]
@@ -743,6 +751,30 @@ class MissingEnvironmentVariable(ApplicationError):
         self.name = name
 
 
+class NetworkPlatformSettings:
+    """Settings required for provisioning a network platform."""
+    def __init__(self, collection, inventory_vars):  # type: (str, t.Type[str, str]) -> None
+        self.collection = collection
+        self.inventory_vars = inventory_vars
+
+
+def get_network_settings(args, platform, version):  # type: (NetworkIntegrationConfig, str, str) -> NetworkPlatformSettings
+    """Returns settings for the given network platform and version."""
+    platform_version = '%s/%s' % (platform, version)
+    completion = get_network_completion().get(platform_version, {})
+    collection = args.platform_collection.get(platform, completion.get('collection'))
+
+    settings = NetworkPlatformSettings(
+        collection,
+        dict(
+            ansible_connection=args.platform_connection.get(platform, completion.get('connection')),
+            ansible_network_os='%s.%s' % (collection, platform) if collection else platform,
+        )
+    )
+
+    return settings
+
+
 def docker_qualify_image(name):
     """
     :type name: str