# HG changeset patch
# User Raphaël Gomès <rgomes@octobus.net>
# Date 1584632583 -3600
#      Thu Mar 19 16:43:03 2020 +0100
# Node ID 0e1b65425e58490ef351b513943671a3b75a9722
# Parent  4f666dffe0cd6c69cc72c5fcaf924c349eac59b6
Create helpers for adding unclean files in the setup phase of a benchmark

diff --git a/benchmarks/utils/__init__.py b/benchmarks/utils/__init__.py
--- a/benchmarks/utils/__init__.py
+++ b/benchmarks/utils/__init__.py
@@ -280,6 +280,32 @@
     return worker_wrapper
 
 
+def add_ignored_files(f):
+    """
+    During the setup phase of the decorated benchmark, extracts a tarball of
+    ignored files made for this repository to fake a dirty working directory.
+    """
+    @functools.wraps(f)
+    def wrapper(self, *args, **kwargs):
+        self.add_unclean_files('ignored', self.repo_name, self.repo_path)
+        return f(self, *args, **kwargs)
+
+    return wrapper
+
+
+def add_unknown_files(f):
+    """
+    During the setup phase of the decorated benchmark, extracts a tarball of
+    unknown files made for this repository to fake a dirty working directory.
+    """
+    @functools.wraps(f)
+    def wrapper(self, *args, **kwargs):
+        self.add_unclean_files('unknown', self.repo_name, self.repo_path)
+        return f(self, *args, **kwargs)
+
+    return wrapper
+
+
 def not_compatible_with(revset, filter_fn=None):
     """Specifies the revset wherein the command is NOT expected to work.
 
@@ -834,6 +860,28 @@
         # in subclasses to mirror super().setup().
         pass
 
+    def add_unclean_files(self, kind, repo_name, root):
+        tarballs = TARBALLS_OF_UNCLEANS[repo_name]
+
+        tarball_name = tarballs.get(kind)
+        if tarball_name is None:
+            # there may not be a tarball yet, just return
+            return
+        tarball = os.path.join(REPOS_DIR,
+                                   'uncleans-tarballs',
+                               tarball_name)
+
+        with open(self.need_update_marker_path, "w"):
+            pass
+
+        try:
+            subprocess.check_call(
+                ["tar", "xf", tarball, '-C', root]
+            )
+        except subprocess.CalledProcessError:
+            msg = "Extracting failed for `%s` tarball (`%s`) for repo `%s`"
+            raise RuntimeError(msg % (kind, tarball, repo_name))
+
 
 class BaseNChangesetsTestSuite(BaseTestSuite):