import argparse import hashlib import io import json import unittest from unittest import mock import run_packed_backup_export as controller from packed_backup_workflow_spec import parse_workflow_spec def cluster_id_object(): content = b"cluster-1\n" return controller.ClusterIdObject( bucket="dfs-bucket", key="packed-backup-csv-export/cluster-id-lists/run-1/cluster-ids", region="us-west-2", sha256=hashlib.sha256(content).hexdigest(), byte_count=len(content), cluster_count=1, ) def workflow_spec(): return parse_workflow_spec( { "schemaVersion": 1, "mysql": {"secretName": "mysql-secret"}, "pack": { "shards": 1, "cseCtlImage": "pack-image", "resources": {"cpuRequest": "1", "memoryRequest": "1Gi"}, }, "export": { "outPrefix": "s3://bucket/export", "cseCtlImage": "export-image", "dumplingImage": "dumpling-image", "resources": {"cpuRequest": "1", "memoryRequest": "1Gi"}, }, "watermark": { "image": "watermark-image", "awsRegion": "us-east-1", "resources": {"cpuRequest": "1", "memoryRequest": "1Gi"}, }, } ) class ControllerTests(unittest.TestCase): def test_job_ttl_is_not_a_top_level_cli_option(self): required = [ "--name", "run-1", "--spec", "/workflow.yaml", "--mysql-defaults-file", "/mysql.cnf", "--local-gpg", "/public.asc", ] parser = controller.build_parser() args = parser.parse_args(required) self.assertFalse(hasattr(args, "ttl_seconds_after_finished")) with mock.patch("sys.stderr", new=io.StringIO()): with self.assertRaises(SystemExit): parser.parse_args( [ *required, "--ttl-seconds-after-finished", "1", ] ) def test_workflow_yaml_ttl_is_applied_to_generated_jobs(self): manifest = { "apiVersion": "v1", "kind": "List", "items": [ {"apiVersion": "v1", "kind": "ConfigMap", "metadata": {}}, { "apiVersion": "batch/v1", "kind": "Job", "metadata": {"name": "job-1"}, "spec": {"ttlSecondsAfterFinished": 1}, }, ], } rendered = json.loads( controller.job_manifest_with_ttl( json.dumps(manifest), "Export", 86400, ) ) self.assertNotIn("spec", rendered["items"][0]) self.assertEqual( rendered["items"][1]["spec"]["ttlSecondsAfterFinished"], 86400, ) def test_generator_commands_preserve_each_argv_boundary(self): payload = workflow_spec().as_dict() payload["pack"]["generatorArgs"] = ["--cpu.limit", "2"] payload["pack"]["commandArgs"] = ["--compression", "zstd"] payload["export"]["generatorArgs"] = ["--max-retries", "3"] payload["export"]["workerArgs"] = ["--parallel", "2"] payload["export"]["dumplingArgs"] = ["--threads", "2"] payload["watermark"]["generatorArgs"] = ["--future-generator", "value"] payload["watermark"]["workerArgs"] = ["--future-worker", "value"] spec = parse_workflow_spec(payload) tasks = [ { "task_name": "task-1", "pack_backup_prefix": "s3://bucket/packed/task-1/", } ] with mock.patch.object( controller, "run_command", return_value="{}", ) as run: controller.generate_pack_manifest( spec, "bash", "run-1", "pack-run", "20260731/000000.meta", "selection", ) pack_command = run.call_args.args[0] self.assertEqual(pack_command.count("--"), 1) self.assertEqual( pack_command[-4:], ["--backup-name", "20260731/000000.meta", "--compression", "zstd"], ) self.assertLess(pack_command.index("--cpu.limit"), pack_command.index("--")) self.assertIn("--cluster-id-dfs-object-key", pack_command) self.assertNotIn("--cluster-id-config-map", pack_command) self.assertEqual( pack_command[pack_command.index("--node-group") + 1], "general", ) self.assertNotIn("--ttl-seconds-after-finished", pack_command) controller.generate_export_manifest( argparse.Namespace(local_gpg="/gpg/public.asc"), spec, "run-1", "export-run", tasks, ) export_command = run.call_args.args[0] self.assertEqual( controller.GENERATE_EXPORT.name, "generate_dumpling_packed_backup_export.py", ) separators = [ index for index, value in enumerate(export_command) if value == "--" ] self.assertEqual(len(separators), 2) self.assertLess(export_command.index("--max-retries"), separators[0]) self.assertLess(export_command.index("--parallel"), separators[1]) self.assertGreater(export_command.index("--threads"), separators[1]) self.assertEqual( export_command[export_command.index("--node-group") + 1], "general", ) self.assertNotIn("--ttl-seconds-after-finished", export_command) controller.generate_watermark_manifest( spec, "run-1", "watermark-run", "export-run", tasks, controller.backup_cutoff_from_name("20260731/000000.meta"), ) watermark_command = run.call_args.args[0] self.assertEqual( controller.GENERATE_WATERMARK.name, "generate_packed_backup_watermark.py", ) self.assertEqual(watermark_command.count("--"), 1) separator = watermark_command.index("--") self.assertLess(watermark_command.index("--future-generator"), separator) self.assertGreater(watermark_command.index("--future-worker"), separator) self.assertNotIn("--cluster-id-object-uri", watermark_command) self.assertNotIn("--cluster-id-config-map", watermark_command) self.assertEqual( watermark_command[watermark_command.index("--node-group") + 1], "general", ) self.assertNotIn("--ttl-seconds-after-finished", watermark_command) def test_workflow_record_manifest_saves_role_and_object_metadata_only(self): selection = cluster_id_object() manifest = json.loads( controller.workflow_record_manifest( "run-1", "pack-run-1", selection, "20260731/000000.meta", ) ) self.assertEqual( manifest["metadata"]["labels"][controller.RESOURCE_ROLE_LABEL], "cluster-selection", ) self.assertEqual( manifest["data"], { "backup-name": "20260731/000000.meta", "cluster-id-count": "1", "cluster-id-object-bytes": "10", "cluster-id-object-region": "us-west-2", "cluster-id-object-sha256": selection.sha256, "cluster-id-object-uri": selection.uri, }, ) self.assertNotIn("cluster-ids", manifest["data"]) def test_workflow_record_ignores_obsolete_effective_config_data(self): manifest = json.loads( controller.workflow_record_manifest( "run-1", "pack-run-1", cluster_id_object(), "20260731/000000.meta", ) ) manifest["data"]["effective-configuration.json"] = "not-json" manifest["data"]["effective-configuration-sha256"] = "not-a-hash" with mock.patch.object( controller, "run_command", return_value=json.dumps(manifest), ): record = controller.load_workflow_record( "kubectl", "run-1", "pack-run-1", ) self.assertIsNotNone(record) assert record is not None self.assertEqual(record.backup_name, "20260731/000000.meta") self.assertEqual(record.cluster_id_object, cluster_id_object()) def test_dfs_storage_config_uses_bucket_and_region_but_not_prefix(self): stateful_set = { "spec": { "template": { "spec": { "containers": [ { "name": "tikv-api", "env": [ {"name": "DFS_PREFIX", "value": "root/prefix"}, {"name": "DFS_S3_BUCKET", "value": "dfs-bucket"}, {"name": "DFS_S3_REGION", "value": "us-west-2"}, ], } ] } } } } with mock.patch.object( controller, "run_command", return_value=json.dumps(stateful_set), ): storage = controller.load_dfs_storage_config("kubectl") self.assertEqual( storage, controller.DfsStorageConfig("dfs-bucket", "us-west-2"), ) def test_cluster_id_object_upload_uses_raw_key_and_conditional_create(self): client = mock.Mock() storage = controller.DfsStorageConfig("dfs-bucket", "us-west-2") selection = controller.upload_cluster_id_object( client, storage, "run-1", ["cluster-1", "cluster-2"], "20260731/000000.meta", ) expected_key = ( "packed-backup-csv-export/cluster-id-lists/run-1/cluster-ids" ) self.assertEqual(selection.key, expected_key) request = client.put_object.call_args.kwargs self.assertEqual(request["Bucket"], "dfs-bucket") self.assertEqual(request["Key"], expected_key) self.assertNotIn("root", request["Key"]) self.assertEqual(request["Body"], b"cluster-1\ncluster-2\n") self.assertEqual(request["IfNoneMatch"], "*") self.assertEqual(request["Metadata"]["sha256"], selection.sha256) def test_recovered_s3_object_content_and_metadata_are_verified(self): content = b" cluster-1 \n\n" response = { "Body": io.BytesIO(content), "ContentLength": len(content), "Metadata": { "workflow-name": "run-1", "backup-name": "20260731/000000.meta", "sha256": hashlib.sha256(content).hexdigest(), "byte-count": str(len(content)), "cluster-count": "1", }, } selection, backup_name = controller.cluster_id_object_from_s3_response( controller.DfsStorageConfig("dfs-bucket", "us-west-2"), "run-1", "raw/object-key", response, ) self.assertEqual(selection.key, "raw/object-key") self.assertEqual(selection.cluster_count, 1) self.assertEqual(backup_name, "20260731/000000.meta") def test_existing_workflow_reuses_frozen_object_without_mysql_or_upload(self): selection = cluster_id_object() record = controller.WorkflowRecord( name="pack-run-1-cluster-ids", cluster_id_object=selection, backup_name="20260731/000000.meta", ) args = argparse.Namespace( backup_from_prefix="20260731", mysql_defaults_file="/mysql.cnf", ) result = {} with mock.patch.object( controller, "load_dfs_storage_config", return_value=controller.DfsStorageConfig( selection.bucket, selection.region, ), ), mock.patch.object( controller, "create_s3_client", return_value=mock.sentinel.s3, ), mock.patch.object( controller, "load_workflow_record", return_value=record, ), mock.patch.object( controller, "validate_cluster_id_object", ) as validate, mock.patch.object( controller, "open_mysql_connection", ) as mysql, mock.patch.object( controller, "upload_cluster_id_object", ) as upload: workflow = controller.ensure_workflow_record( args, {"kubectl": "kubectl"}, "run-1", "pack-run-1", result, ) self.assertEqual(workflow.record, record) self.assertEqual(result["selected_clusters"], 1) validate.assert_called_once_with( mock.sentinel.s3, "run-1", selection, "20260731/000000.meta", ) mysql.assert_not_called() upload.assert_not_called() def test_object_created_before_configmap_is_recovered_without_reselection(self): selection = cluster_id_object() recovered_record = controller.WorkflowRecord( name="pack-run-1-cluster-ids", cluster_id_object=selection, backup_name="20260731/000000.meta", ) args = argparse.Namespace( backup_from_prefix="changed-prefix", mysql_defaults_file="/changed.cnf", ) result = {} with mock.patch.object( controller, "load_dfs_storage_config", return_value=controller.DfsStorageConfig( selection.bucket, selection.region, ), ), mock.patch.object( controller, "create_s3_client", return_value=mock.sentinel.s3, ), mock.patch.object( controller, "load_workflow_record", return_value=None, ), mock.patch.object( controller, "list_instance_jobs", return_value=[], ), mock.patch.object( controller, "load_existing_cluster_id_object", return_value=(selection, "20260731/000000.meta"), ), mock.patch.object( controller, "create_workflow_record", return_value=recovered_record, ) as create_record, mock.patch.object( controller, "discover_backup", ) as discover, mock.patch.object( controller, "open_mysql_connection", ) as mysql: workflow = controller.ensure_workflow_record( args, {"kubectl": "kubectl"}, "run-1", "pack-run-1", result, ) self.assertEqual(workflow.record, recovered_record) self.assertTrue(result["workflow_record_created"]) create_record.assert_called_once_with( "kubectl", "run-1", "pack-run-1", selection, "20260731/000000.meta", ) discover.assert_not_called() mysql.assert_not_called() def test_saved_cluster_object_parser_trims_and_ignores_empty_lines(self): self.assertEqual( controller.parse_cluster_id_object_content( b" cluster-1 \n\ncluster-2\r\n" ), ["cluster-1", "cluster-2"], ) def test_cluster_selection_reads_mysql_in_bounded_batches(self): connection = mock.MagicMock() cursor = connection.cursor.return_value.__enter__.return_value cursor.fetchmany.side_effect = [ [("cluster-1",), ("cluster-2",)], [("cluster-3",)], [], ] selected = controller.select_cluster_ids( connection, controller.backup_cutoff_from_name("20260731/000000.meta"), ) self.assertEqual(selected, ["cluster-1", "cluster-2", "cluster-3"]) self.assertEqual(cursor.fetchmany.call_count, 3) cursor.fetchall.assert_not_called() query = cursor.execute.call_args.args[0] self.assertIn("NOT EXISTS", query) self.assertIn(controller.CLUSTER_BLOCKLIST_TABLE, query) self.assertIn( "blocklist.cluster_id = active_clusters.cluster_id", query, ) def test_watermark_result_uses_the_named_terminated_container(self): result = { "status": "succeeded", "watermark_instance": "watermark-run", "export_instance": "export-run", "backup_time": "2026-07-31T00:00:00Z", "packed_clusters": 1, "packed_meta_objects": 2, "verified_complete_markers": 2, "unverified_complete_marker_clusters": 0, "mysql_updated_clusters": 1, "newer_complete_marker_clusters": 0, "newer_mysql_watermark_clusters": 0, "mysql_update_failed_clusters": 0, } pod_list = { "items": [ { "status": { "containerStatuses": [ { "name": "sidecar", "state": { "terminated": {"message": "ignored"} }, }, { "name": controller.WATERMARK_CONTAINER, "state": { "terminated": { "message": json.dumps(result) } }, }, ] } } ] } with mock.patch.object( controller, "run_command", return_value=json.dumps(pod_list), ) as run: actual = controller.load_watermark_job_result( "kubectl", "watermark-run", "export-run", controller.backup_cutoff_from_name("20260731/000000.meta"), ) self.assertEqual(actual["status"], "succeeded") self.assertEqual(actual["verified_complete_markers"], 2) command = run.call_args.args[0] self.assertEqual(command[4:6], ["get", "pods"]) self.assertNotIn("logs", command) self.assertIn( f"{controller.INSTANCE_LABEL}=watermark-run", command, ) def test_watermark_result_requires_termination_message(self): pod_list = { "items": [ { "status": { "containerStatuses": [ { "name": controller.WATERMARK_CONTAINER, "state": {"terminated": {"exitCode": 1}}, } ] } } ] } with mock.patch.object( controller, "run_command", return_value=json.dumps(pod_list), ), self.assertRaisesRegex( controller.WorkflowError, "no termination message", ): controller.load_watermark_job_result( "kubectl", "watermark-run", "export-run", controller.backup_cutoff_from_name("20260731/000000.meta"), ) def test_wait_for_jobs_has_no_overall_deadline(self): active = { "metadata": {"name": "job-1"}, "status": {"conditions": []}, } complete = { "metadata": {"name": "job-1"}, "status": {"conditions": [{"type": "Complete", "status": "True"}]}, } with mock.patch.object( controller, "list_instance_jobs", side_effect=[[active], [complete]], ), mock.patch.object(controller.time, "sleep") as sleep: controller.wait_for_jobs( "kubectl", "instance", ["job-1"], "Pack", ) sleep.assert_called_once_with(controller.JOB_POLL_INTERVAL_SECONDS) if __name__ == "__main__": unittest.main()