Fix --save-location bugs for collect command

This addresses 2 concerns:

1) When -s or --save-location is omitted for pegleg
   site collect command, the following error is raised:

     TypeError: '_io.TextIOWrapper' object is not subscriptable

   That is because the default kwarg for the associated option
   is sys.stdout [0] which returns an object not a path, causing
   the Click library itself to blow up. Note that this is dropped
   in master [1], which isn't correct, as there is a desire to
   use stdout as the default location for collection output.

   Thus, new logic is added to output the collected document data
   to stdout.

2) When a directory doesn't exist and -s is provided (as an example
   pegleg site -p /path/to/primary/dir collect <site_name> -s <non_existent_dir>)
   it is better to just create it via os.makedirs() than to error out
   with an exception [2].

[0] d9692126ed (diff-bd81bbb896486546c9e59d018a83d05d)
[1] f25b5b6593/src/bin/pegleg/pegleg/cli.py (L72)
[2] 843d1a5010/src/bin/pegleg/pegleg/engine/site.py (L36)

Change-Id: Ie1da73efa437cbdfbe6d2b7ab616a1467eb57358

src/bin/pegleg/pegleg/cli.py

@@ -68,9 +68,8 @@ def site(primary_repo, aux_repo):
     '-s',
     '--save-location',
     'save_location',
-    type=click.Path(
+    help='Directory to output the complete site definition. Created '
-        file_okay=False, dir_okay=True, writable=True, resolve_path=True),
+    'automatically if it does not already exist.')
-    help='Where to output the complete site definition.')
 @click.option(
     '--validate',
     'validate',
@@ -99,6 +98,9 @@ def collect(*, save_location, validate, exclude_lint, warn_lint, site_name):
     defines the entire site definition and contains all documents required
     for ingestion by Airship.
 
+    If ``save_location`` isn't specified, then the output is directed to
+    stdout.
+
     Collect can lint documents prior to collection if the ``--validate``
     flag is optionally included.
     """

src/bin/pegleg/pegleg/engine/site.py

@@ -27,14 +27,45 @@ __all__ = ['collect', 'impacted', 'list_', 'show', 'render']
 LOG = logging.getLogger(__name__)
 
 
-def collect(site_name, save_location):
+def _read_and_format_yaml(filename):
-    try:
+    with open(filename) as f:
-        save_files = dict()
+        lines_to_write = f.readlines()
-        if save_location is None:
+        if lines_to_write[0] != '---\n':
-            raise ValueError('Missing param: save-location')
+            lines_to_write = ['---\n'] + lines_to_write
-        elif not os.path.exists(save_location):
+        if lines_to_write[-1] != '...\n':
-            raise FileNotFoundError('Invalid save-location path')
+            lines_to_write.append('...\n')
+    return lines_to_write or []
 
+
+def _collect_to_stdout(site_name):
+    """Collects all documents related to ``site_name`` and outputs them to
+    stdout via ``output_stream``.
+    """
+    try:
+        for repo_base, filename in util.definition.site_files_by_repo(
+                site_name):
+            for line in _read_and_format_yaml(filename):
+                # This code is a pattern to convert \r\n to \n.
+                click.echo("\n".join(line.splitlines()))
+    except Exception as ex:
+        raise click.ClickException("Error printing output: %s" % str(ex))
+
+
+def _collect_to_file(site_name, save_location):
+    """Collects all documents related to ``site_name`` and outputs them to
+    the file denoted by ``save_location``.
+    """
+    if not os.path.exists(save_location):
+        LOG.debug("Collection save location %s does not exist. Creating "
+                  "automatically.", save_location)
+        os.makedirs(save_location)
+    # In case save_location already exists and isn't a directory.
+    if not os.path.isdir(save_location):
+        raise click.ClickException('save_location %s already exists, but must '
+                                   'be a directory')
+
+    save_files = dict()
+    try:
         for repo_base, filename in util.definition.site_files_by_repo(
                 site_name):
             repo_name = os.path.normpath(repo_base).split(os.sep)[-1]
@@ -42,14 +73,7 @@ def collect(site_name, save_location):
             if repo_name not in save_files:
                 save_files[repo_name] = open(save_file, "w")
             LOG.debug("Collecting file %s to file %s" % (filename, save_file))
-
+            save_files[repo_name].writelines(_read_and_format_yaml(filename))
-            with open(filename) as f:
-                lines_to_write = f.readlines()
-                if lines_to_write[0] != '---\n':
-                    lines_to_write = ['---\n'] + lines_to_write
-                if lines_to_write[-1] != '...\n':
-                    lines_to_write.append('...\n')
-                save_files[repo_name].writelines(lines_to_write)
     except Exception as ex:
         raise click.ClickException("Error saving output: %s" % str(ex))
     finally:
@@ -57,6 +81,13 @@ def collect(site_name, save_location):
             f.close()
 
 
+def collect(site_name, save_location):
+    if save_location:
+        _collect_to_file(site_name, save_location)
+    else:
+        _collect_to_stdout(site_name)
+
+
 def impacted(input_stream, output_stream):
     mapping = _build_impact_mapping()
     impacted_sites = set()

src/bin/pegleg/tests/unit/engine/test_collect.py

@@ -13,8 +13,12 @@
 # limitations under the License.
 
 import mock
+import os
+import shutil
 import yaml
 
+import click
+
 from pegleg import cli
 from pegleg import config
 from pegleg.engine import errorcodes
@@ -25,8 +29,8 @@ from pegleg.engine.util import files
 from tests.unit.fixtures import create_tmp_deployment_files
 
 
-def _test_site_collect(expected_site, collection_path):
+def _site_definition(site_name):
-    site_definition = {
+    SITE_DEFINITION = {
         'schema': 'pegleg/SiteDefinition/v1',
         'metadata': {
             'storagePolicy': 'cleartext',
@@ -35,39 +39,87 @@ def _test_site_collect(expected_site, collection_path):
                 'abstract': False,
                 'layer': 'site'
             },
-            'name': expected_site
+            'name': site_name
         },
         'data': {
-            'site_type': expected_site,
+            'site_type': site_name,
             'revision': 'v1.0'
         }
     }
-    expected_document_names = [
+    return SITE_DEFINITION
+
+
+def _expected_document_names(site_name):
+    EXPECTED_DOCUMENT_NAMES = [
         'global-common',
         'global-v1.0',
-        '%s-type-common' % expected_site,
+        '%s-type-common' % site_name,
-        '%s-type-v1.0' % expected_site,
+        '%s-type-v1.0' % site_name,
-        site_definition['metadata']['name'],
+        _site_definition(site_name)["metadata"]["name"],
-        '%s-chart' % expected_site,
+        '%s-chart' % site_name,
-        '%s-passphrase' % expected_site,
+        '%s-passphrase' % site_name,
     ]
+    return EXPECTED_DOCUMENT_NAMES
 
-    deployment_files = collection_path.join('deployment_files.yaml')
-    assert deployment_files.isfile()
 
-    with open(str(deployment_files), 'r') as f:
+def _test_site_collect_to_file(tmpdir, site_name, collection_path):
-        lines = f.read()
+    try:
+        collection_path = tmpdir.mkdir("cicd_path")
+        collection_str_path = str(collection_path)
+        site.collect(site_name, collection_str_path)
+
+        deployment_files = collection_path.join('deployment_files.yaml')
+        assert deployment_files.isfile()
+
+        with open(str(deployment_files), 'r') as f:
+            lines = f.read()
         deployment_documents = list(yaml.safe_load_all(lines))
 
-    assert sorted(expected_document_names) == sorted(
+        assert sorted(_expected_document_names(site_name)) == sorted(
-        [x['metadata']['name'] for x in deployment_documents])
+            [x['metadata']['name'] for x in deployment_documents])
+    finally:
+        if os.path.exists(collection_str_path):
+            shutil.rmtree(collection_str_path, ignore_errors=True)
 
 
-def test_site_collect(tmpdir, create_tmp_deployment_files):
+def test_site_collect_to_file(tmpdir, create_tmp_deployment_files):
-    collection_path = tmpdir.mkdir("cicd_path")
+    _test_site_collect_to_file(tmpdir, "cicd", "cicd_path")
-    site.collect("cicd", str(collection_path))
+    _test_site_collect_to_file(tmpdir, "lab", "lab_path")
-    _test_site_collect("cicd", collection_path)
 
-    collection_path = tmpdir.mkdir("lab_path")
+
-    site.collect("lab", str(collection_path))
+def _test_site_collect_to_stdout(site_name):
-    _test_site_collect("lab", collection_path)
+    # 2nd arg of None will force redirection to stdout.
+    with mock.patch.object(click, 'echo') as mock_echo:
+        site.collect(site_name, None)
+
+    expected_names = _expected_document_names(site_name)
+    all_lines = [x[1][0].strip() for x in mock_echo.mock_calls]
+
+    assert all_lines, "Nothing written to stdout"
+    for expected in expected_names:
+        assert 'name: %s' % expected in all_lines
+
+
+def test_site_collect_to_stdout(create_tmp_deployment_files):
+    _test_site_collect_to_stdout("cicd")
+    _test_site_collect_to_stdout("lab")
+
+
+def test_read_and_format_yaml(tmpdir):
+    # Validate the case where the YAML already begins with leading --- and ends
+    # with trailing ... -- there should be no change.
+    tempdir = tmpdir.mkdir(__name__)
+    tempfile = tempdir.join("test-manifests-1.yaml")
+    # Check that Linux-style newline is output instead when \r\n is provided.
+    tempfile.write("---\r\nfoo:bar\r\n...\r\n")
+    output = list(site._read_and_format_yaml(str(tempfile)))
+    expected = ['---\n', 'foo:bar\n', '...\n']
+    assert expected == output
+
+    # Validate the case where the YAML is missing --- & ... and check that
+    # they are added.
+    tempfile = tempdir.join("test-manifests-2.yaml")
+    tempfile.write("foo:bar\n")
+    output = list(site._read_and_format_yaml(str(tempfile)))
+    expected = ['---\n', 'foo:bar\n', '...\n']
+    assert expected == output