more work on resuming.

2025-12-10 21:07:52 -08:00
parent c3d31b4ab5
commit 6b4f3939a5
5 changed files with 512 additions and 532 deletions
--- a/test/wikiq_test_utils.py
+++ b/test/wikiq_test_utils.py
@@ -0,0 +1,75 @@
+import os
+import shutil
+import subprocess
+from typing import Final, Union
+
+TEST_DIR: Final[str] = os.path.dirname(os.path.realpath(__file__))
+WIKIQ: Final[str] = os.path.join(os.path.join(TEST_DIR, ".."), "src/wikiq/__init__.py")
+TEST_OUTPUT_DIR: Final[str] = os.path.join(TEST_DIR, "test_output")
+BASELINE_DIR: Final[str] = os.path.join(TEST_DIR, "baseline_output")
+
+IKWIKI: Final[str] = "ikwiki-20180301-pages-meta-history"
+SAILORMOON: Final[str] = "sailormoon"
+TWINPEAKS: Final[str] = "twinpeaks"
+REGEXTEST: Final[str] = "regextest"
+
+
+class WikiqTester:
+    def __init__(
+        self,
+        wiki: str,
+        case_name: str,
+        suffix: Union[str, None] = None,
+        in_compression: str = "bz2",
+        baseline_format: str = "tsv",
+        out_format: str = "tsv",
+    ):
+        self.input_file = os.path.join(
+            TEST_DIR, "dumps", "{0}.xml.{1}".format(wiki, in_compression)
+        )
+
+        basename = "{0}_{1}".format(case_name, wiki)
+        if suffix:
+            basename = "{0}_{1}".format(basename, suffix)
+
+        self.output = os.path.join(
+            TEST_OUTPUT_DIR, "{0}.{1}".format(basename, out_format)
+        )
+
+        if os.path.exists(self.output):
+            if os.path.isfile(self.output):
+                os.remove(self.output)
+            else:
+                shutil.rmtree(self.output)
+
+        if out_format == "parquet":
+            os.makedirs(self.output, exist_ok=True)
+
+        if suffix is None:
+            self.wikiq_baseline_name = "{0}.{1}".format(wiki, baseline_format)
+            self.wikiq_out_name = "{0}.{1}".format(wiki, out_format)
+        else:
+            self.wikiq_baseline_name = "{0}_{1}.{2}".format(
+                wiki, suffix, baseline_format
+            )
+            self.wikiq_out_name = "{0}_{1}.{2}".format(wiki, suffix, out_format)
+
+        if case_name is not None:
+            self.baseline_file = os.path.join(
+                BASELINE_DIR, "{0}_{1}".format(case_name, self.wikiq_baseline_name)
+            )
+
+    def call_wikiq(self, *args: str, out: bool = True):
+        """
+        Calls wikiq with the passed arguments on the input file relevant to the test.
+        :param args: The command line arguments to pass to wikiq.
+        :param out: Whether to pass an output argument to wikiq.
+        :return: The output of the wikiq call.
+        """
+        if out:
+            call = " ".join([WIKIQ, self.input_file, "-o", self.output, "--batch-size", "10", *args])
+        else:
+            call = " ".join([WIKIQ, self.input_file, "--batch-size", "10", *args])
+
+        print(call)
+        return subprocess.check_output(call, stderr=subprocess.PIPE, shell=True)