element_generators.py: Add a utility function to create arbitrary size elements

New file consisting of functions that allow us to generate elements
on the fly
__init__.py now includes create_element_size()

2 files changed, 41 insertions, 0 deletions

diff --git a/tests/testutils/__init__.py b/tests/testutils/__init__.py
index 7e5b792a2..93143b505 100644
--- a/tests/testutils/__init__.py
+++ b/tests/testutils/__init__.py
@@ -1,3 +1,4 @@
 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
diff --git a/tests/testutils/element_generators.py b/tests/testutils/element_generators.py
new file mode 100644
index 000000000..3f6090da8
--- /dev/null
+++ b/tests/testutils/element_generators.py
@@ -0,0 +1,40 @@
+import os
+
+from buildstream import _yaml
+
+
+# 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
+#
+# 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
+#
+# Returns:
+#  Nothing (creates a .bst file of specified size)
+#
+def create_element_size(name, path, dependencies, size):
+    os.makedirs(path, exist_ok=True)
+
+    # Create a file to be included in this element's artifact
+    with open(os.path.join(path, name + '_data'), 'wb+') as f:
+        f.write(os.urandom(size))
+
+    # Simplest case: We want this file (of specified size) to just
+    # be an import element.
+    element = {
+        'kind': 'import',
+        'sources': [
+            {
+                'kind': 'local',
+                'path': os.path.join(path, name + '_data')
+            }
+        ],
+        'depends': dependencies
+    }
+    _yaml.dump(element, os.path.join(path, name))