testutils/element_generators.py: Changing sized element functions

 * create_element_size()

   Now uses a git Repo object instead of a local source, and
   returns the repo.

 * update_element_size()

   Added this function which can now resize the expected output
   of an element generated with create_element_size(), useful
   to allow testing sized elements with the tracking feature.

2 files changed, 77 insertions, 18 deletions

diff --git a/tests/testutils/__init__.py b/tests/testutils/__init__.py
index c2fae1cc4..4a79c3be2 100644
--- a/tests/testutils/__init__.py
+++ b/tests/testutils/__init__.py
@@ -26,6 +26,6 @@
 from .runcli import cli, cli_integration
 from .repo import create_repo, ALL_REPO_KINDS
 from .artifactshare import create_artifact_share
-from .element_generators import create_element_size
+from .element_generators import create_element_size, update_element_size
 from .junction import generate_junction
 from .runner_integration import wait_for_cache_granularity
diff --git a/tests/testutils/element_generators.py b/tests/testutils/element_generators.py
index 49f235c61..448c8571a 100644
--- a/tests/testutils/element_generators.py
+++ b/tests/testutils/element_generators.py
@@ -1,41 +1,100 @@
 import os
 
 from buildstream import _yaml
+from buildstream import utils
+
+from . import create_repo
 
 
 # create_element_size()
 #
-# This will open a "<name>_data" file for writing and write
-# <size> MB of urandom (/dev/urandom) "stuff" into the file.
-# A bst import element file is then created: <name>.bst
+# Creates an import element with a git repo, using random
+# data to create a file in that repo of the specified size,
+# such that building it will add an artifact of the specified
+# size to the artifact cache.
 #
 # Args:
-#  name: (str) of the element name (e.g. target.bst)
-#  path: (str) pathway to the project/elements directory
-#  dependencies: A list of strings (can also be an empty list)
-#  size: (int) size of the element in bytes
+#    name: (str) of the element name (e.g. target.bst)
+#    project_dir (str): The path to the project
+#    element_path (str): The element path within the project
+#    dependencies: A list of strings (can also be an empty list)
+#    size: (int) size of the element in bytes
 #
 # Returns:
-#  Nothing (creates a .bst file of specified size)
+#    (Repo): A git repo which can be used to introduce trackable changes
+#            by using the update_element_size() function below.
 #
 def create_element_size(name, project_dir, elements_path, dependencies, size):
     full_elements_path = os.path.join(project_dir, elements_path)
     os.makedirs(full_elements_path, exist_ok=True)
 
-    # Create a file to be included in this element's artifact
-    with open(os.path.join(project_dir, name + '_data'), 'wb+') as f:
-        f.write(os.urandom(size))
+    # Create a git repo
+    repodir = os.path.join(project_dir, 'repos')
+    repo = create_repo('git', repodir, subdir=name)
+
+    with utils._tempdir(dir=project_dir) as tmp:
+
+        # We use a data/ subdir in the git repo we create,
+        # and we set the import element to only extract that
+        # part; this ensures we never include a .git/ directory
+        # in the cached artifacts for these sized elements.
+        #
+        datadir = os.path.join(tmp, 'data')
+        os.makedirs(datadir)
+
+        # Use /dev/urandom to create the sized file in the datadir
+        with open(os.path.join(datadir, name), 'wb+') as f:
+            f.write(os.urandom(size))
+
+        # Create the git repo from the temp directory
+        ref = repo.create(tmp)
 
-    # Simplest case: We want this file (of specified size) to just
-    # be an import element.
     element = {
         'kind': 'import',
         'sources': [
-            {
-                'kind': 'local',
-                'path': name + '_data'
-            }
+            repo.source_config(ref=ref)
         ],
+        'config': {
+            # Extract only the data directory
+            'source': 'data'
+        },
         'depends': dependencies
     }
     _yaml.dump(element, os.path.join(project_dir, elements_path, name))
+
+    # Return the repo, so that it can later be used to add commits
+    return repo
+
+
+# update_element_size()
+#
+# Updates a repo returned by create_element_size() such that
+# the newly added commit is completely changed, and has the newly
+# specified size.
+#
+# The name and project_dir arguments must match the arguments
+# previously given to create_element_size()
+#
+# Args:
+#    name: (str) of the element name (e.g. target.bst)
+#    project_dir (str): The path to the project
+#    repo: (Repo) The Repo returned by create_element_size()
+#    size: (int) The new size which the element generates, in bytes
+#
+# Returns:
+#    (Repo): A git repo which can be used to introduce trackable changes
+#            by using the update_element_size() function below.
+#
+def update_element_size(name, project_dir, repo, size):
+
+    with utils._tempdir(dir=project_dir) as tmp:
+
+        new_file = os.path.join(tmp, name)
+
+        # Use /dev/urandom to create the sized file in the datadir
+        with open(new_file, 'wb+') as f:
+            f.write(os.urandom(size))
+
+        # Modify the git repo with a new commit to the same path,
+        # replacing the original file with a new one.
+        repo.modify_file(new_file, os.path.join('data', name))