import json import os import shutil import stat import subprocess import tempfile import unittest from pathlib import Path import generate_dumpling_packed_backup_export as export_generator import generate_packed_backup_watermark as watermark_generator ROOT = Path(__file__).resolve().parents[1] WORKFLOW_LABEL = "cse.tidbcloud.com/workflow-name" STAGE_LABEL = "cse.tidbcloud.com/stage" JOB_LABEL = "cse.tidbcloud.com/job-name" ROLE_LABEL = "cse.tidbcloud.com/resource-role" class PythonGeneratorTests(unittest.TestCase): def test_export_manifest_has_trace_labels_and_no_deadline(self): args = export_generator.build_parser().parse_args( [ "--name", "export-run", "--workflow-name", "run-1", "--batch-task-json=-", ] ) starter_args = [ "--out-prefix", "s3://bucket/export", "--flat", "--", "--threads", "2", ] export_generator.validate_options(args, starter_args) manifest = export_generator.build_manifest( args, starter_args, [ { "task_name": "pack-task", "pack_backup_prefix": "s3://bucket/packed/task/", } ], [], "print('init')\n", b"public-key", ) config_map, job = manifest["items"] self.assertEqual(config_map["metadata"]["labels"][WORKFLOW_LABEL], "run-1") self.assertEqual(config_map["metadata"]["labels"][STAGE_LABEL], "export") self.assertEqual(config_map["metadata"]["labels"][ROLE_LABEL], "export-init") self.assertIn("dumpling_packed_backup_init.py", config_map["data"]) self.assertNotIn(JOB_LABEL, config_map["metadata"]["labels"]) self.assertEqual(job["metadata"]["labels"][JOB_LABEL], job["metadata"]["name"]) self.assertEqual( job["spec"]["template"]["metadata"]["labels"], job["metadata"]["labels"], ) self.assertEqual( job["spec"]["ttlSecondsAfterFinished"], 30 * 24 * 60 * 60, ) container = job["spec"]["template"]["spec"]["containers"][0] self.assertEqual( container["command"], [ export_generator.RUNNER_PYTHON, "/usr/local/bin/dumpling_packed_backup_init.py", ], ) init_mount = next( mount for mount in container["volumeMounts"] if mount["name"] == "init-script" ) self.assertEqual( init_mount["subPath"], "dumpling_packed_backup_init.py", ) self.assertNotIn("activeDeadlineSeconds", job["spec"]) def test_watermark_manifest_forwards_worker_argv_and_has_trace_labels(self): base_generator_args = [ "--name", "watermark-run", "--workflow-name", "run-1", "--tasks-json=-", "--mysql-secret", "mysql-secret", "--out-prefix", "s3://bucket/export", "--export-instance", "export-run", "--backup-time", "2026-07-31T00:00:00Z", "--aws-region", "us-east-1", "--runner-image", "runner-image", ] worker_args = ["--future-worker-option", "value"] args = watermark_generator.build_parser().parse_args( base_generator_args ) watermark_generator.validate_options(args, worker_args) manifest = watermark_generator.build_manifest( args, [{"pack_backup_prefix": "s3://bucket/packed/task/"}], "print('worker')\n", worker_args, ) config_map, job = manifest["items"] self.assertEqual( config_map["metadata"]["labels"][ROLE_LABEL], "watermark-input", ) self.assertIn("packed_backup_watermark.py", config_map["data"]) for resource in (config_map, job): self.assertEqual(resource["metadata"]["labels"][WORKFLOW_LABEL], "run-1") self.assertEqual(resource["metadata"]["labels"][STAGE_LABEL], "watermark") self.assertEqual(job["metadata"]["labels"][JOB_LABEL], "watermark-run") self.assertEqual( job["spec"]["ttlSecondsAfterFinished"], 30 * 24 * 60 * 60, ) self.assertEqual( job["spec"]["template"]["metadata"]["labels"], job["metadata"]["labels"], ) self.assertEqual( job["spec"]["template"]["spec"]["containers"][0]["args"][-2:], worker_args, ) pod_spec = job["spec"]["template"]["spec"] self.assertEqual( pod_spec["containers"][0]["terminationMessagePath"], "/dev/termination-log", ) self.assertEqual( pod_spec["containers"][0]["terminationMessagePolicy"], "File", ) self.assertEqual( pod_spec["containers"][0]["command"], [ "python3", "/var/run/packed-backup-watermark/input/" "packed_backup_watermark.py", ], ) container_args = pod_spec["containers"][0]["args"] self.assertNotIn("--cluster-ids", container_args) self.assertNotIn("--cluster-id-object-uri", container_args) self.assertNotIn( "workflow", [volume["name"] for volume in pod_spec["volumes"]], ) self.assertNotIn( "workflow", [ mount["name"] for mount in pod_spec["containers"][0]["volumeMounts"] ], ) self.assertNotIn("activeDeadlineSeconds", job["spec"]) def test_duplicate_generator_owned_option_is_rejected(self): with self.assertRaisesRegex( export_generator.GeneratorError, "must not be repeated", ): export_generator.reject_duplicate_generator_options( ["--name", "one", "--name=two"] ) def test_export_generator_has_no_job_ttl_cli_option(self): with self.assertRaisesRegex( export_generator.GeneratorError, "unrecognized arguments", ): export_generator.build_parser().parse_args( ["--ttl-seconds-after-finished", "1"] ) class PackGeneratorTests(unittest.TestCase): @unittest.skipUnless(shutil.which("jq"), "jq is required") def test_pack_manifest_labels_jobs_pods_configmap_and_ephemeral_pvcs(self): stateful_set = { "spec": { "template": { "spec": { "containers": [ { "name": "tikv-api", "image": "source-image", "env": [ {"name": "DFS_PREFIX", "value": "root"}, {"name": "DFS_S3_BUCKET", "value": "bucket"}, {"name": "DFS_S3_REGION", "value": "us-east-1"}, {"name": "CSE_MASTER_KEY_ID", "value": "id"}, {"name": "CSE_MASTER_KEY_CIPHER_TEXT", "value": "cipher"}, {"name": "CSE_MASTER_KEY_VENDOR", "value": "aws"}, ], } ] } } } } with tempfile.TemporaryDirectory() as directory: temporary = Path(directory) fake_kubectl = temporary / "kubectl" fake_kubectl.write_text( "#!/usr/bin/env bash\n" "if [[ \" $* \" == *\" get statefulset \"* ]]; then\n" f" printf '%s\\n' {json.dumps(json.dumps(stateful_set))}\n" "else\n" " printf 'unexpected kubectl call: %s\\n' \"$*\" >&2\n" " exit 1\n" "fi\n", encoding="utf-8", ) fake_kubectl.chmod(fake_kubectl.stat().st_mode | stat.S_IXUSR) cluster_ids = temporary / "cluster-ids" cluster_ids.write_text("cluster-1\n", encoding="utf-8") environment = dict(os.environ) environment["PATH"] = f"{temporary}{os.pathsep}{environment['PATH']}" result = subprocess.run( [ "bash", str(ROOT / "generate-pack-backup-offline.sh"), "--name", "pack-run", "--workflow-name", "run-1", "--shards", "2", "--image", "pack-image", "--node-group", "adhoc", "--ebs.request", "10Gi", "--ebs.storage-class", "gp3", "--cluster-id-in-file", str(cluster_ids), "--output", "json", "--", "--backup-name", "20260731/000000.meta", ], env=environment, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False, ) self.assertEqual(result.returncode, 0, result.stderr) manifest = json.loads(result.stdout) config_map = manifest["items"][0] self.assertEqual(config_map["kind"], "ConfigMap") self.assertEqual(config_map["metadata"]["labels"][WORKFLOW_LABEL], "run-1") self.assertEqual(config_map["metadata"]["labels"][ROLE_LABEL], "cluster-selection") for job in manifest["items"][1:]: labels = job["metadata"]["labels"] self.assertEqual(labels[WORKFLOW_LABEL], "run-1") self.assertEqual(labels[STAGE_LABEL], "pack") self.assertEqual(labels[JOB_LABEL], job["metadata"]["name"]) self.assertEqual(job["spec"]["template"]["metadata"]["labels"], labels) self.assertEqual( job["spec"]["ttlSecondsAfterFinished"], 30 * 24 * 60 * 60, ) pod_spec = job["spec"]["template"]["spec"] self.assertEqual( pod_spec["nodeSelector"]["serverless.tidbcloud.com/node"], "adhoc", ) self.assertEqual(pod_spec["tolerations"][0]["value"], "adhoc") pvc_labels = ( pod_spec["volumes"][0] ["ephemeral"]["volumeClaimTemplate"]["metadata"]["labels"] ) self.assertEqual(pvc_labels[WORKFLOW_LABEL], "run-1") self.assertEqual(pvc_labels[STAGE_LABEL], "pack") self.assertEqual(pvc_labels[JOB_LABEL], job["metadata"]["name"]) @unittest.skipUnless(shutil.which("jq"), "jq is required") def test_pack_dfs_object_key_is_raw_and_needs_no_configmap_mount(self): stateful_set = { "spec": { "template": { "spec": { "containers": [ { "name": "tikv-api", "image": "source-image", "env": [ {"name": "DFS_PREFIX", "value": "root"}, {"name": "DFS_S3_BUCKET", "value": "bucket"}, {"name": "DFS_S3_REGION", "value": "us-east-1"}, {"name": "CSE_MASTER_KEY_ID", "value": "id"}, { "name": "CSE_MASTER_KEY_CIPHER_TEXT", "value": "cipher", }, {"name": "CSE_MASTER_KEY_VENDOR", "value": "aws"}, ], } ] } } } } object_key = "selections/run-1/cluster-ids" with tempfile.TemporaryDirectory() as directory: temporary = Path(directory) fake_kubectl = temporary / "kubectl" fake_kubectl.write_text( "#!/usr/bin/env bash\n" f"printf '%s\\n' {json.dumps(json.dumps(stateful_set))}\n", encoding="utf-8", ) fake_kubectl.chmod(fake_kubectl.stat().st_mode | stat.S_IXUSR) environment = dict(os.environ) environment["PATH"] = f"{temporary}{os.pathsep}{environment['PATH']}" result = subprocess.run( [ "bash", str(ROOT / "generate-pack-backup-offline.sh"), "--name", "pack-run", "--workflow-name", "run-1", "--image", "pack-image", "--cluster-id-dfs-object-key", object_key, "--output", "json", "--", "--backup-name", "20260731/000000.meta", ], env=environment, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False, ) self.assertEqual(result.returncode, 0, result.stderr) manifest = json.loads(result.stdout) self.assertEqual([item["kind"] for item in manifest["items"]], ["Job"]) self.assertEqual( manifest["items"][0]["spec"]["ttlSecondsAfterFinished"], 30 * 24 * 60 * 60, ) pod_spec = manifest["items"][0]["spec"]["template"]["spec"] command = pod_spec["containers"][0]["command"] option_index = command.index("--keyspace-select.cluster-id-in") self.assertEqual(command[option_index + 1], f"dfs:{object_key}") self.assertNotEqual(command[option_index + 1], f"dfs:root/{object_key}") self.assertEqual( [volume["name"] for volume in pod_spec["volumes"]], ["data"], ) self.assertEqual( [ mount["name"] for mount in pod_spec["containers"][0]["volumeMounts"] ], ["data"], ) def test_pack_rejects_empty_dfs_object_key(self): result = subprocess.run( [ "bash", str(ROOT / "generate-pack-backup-offline.sh"), "--cluster-id-dfs-object-key=", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False, ) self.assertNotEqual(result.returncode, 0) self.assertIn( "--cluster-id-dfs-object-key requires a value", result.stderr, ) def test_pack_has_no_job_ttl_cli_option(self): result = subprocess.run( [ "bash", str(ROOT / "generate-pack-backup-offline.sh"), "--ttl-seconds-after-finished=1", "--", "--backup-name", "20260731/000000.meta", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False, ) self.assertNotEqual(result.returncode, 0) self.assertIn("unknown generator option", result.stderr) if __name__ == "__main__": unittest.main()