import importlib.util import json import tempfile import unittest from pathlib import Path from types import SimpleNamespace from unittest import mock SCRIPT_PATH = ( Path(__file__).parents[1] / ".github" / "scripts" / "byoc_on_demand.py" ) SPEC = importlib.util.spec_from_file_location("byoc_on_demand", SCRIPT_PATH) MODULE = importlib.util.module_from_spec(SPEC) assert SPEC.loader is not None SPEC.loader.exec_module(MODULE) DEV_REGISTRY = "385595570414.dkr.ecr.us-west-2.amazonaws.com" DIGEST = "sha256:" + "a" * 64 class ByocOnDemandTest(unittest.TestCase): def test_parses_exact_source_without_using_an_image_config(self): images = MODULE.parse_images( json.dumps( [ { "source": f"{DEV_REGISTRY}/tidbcloud/pd:v26.3.4-nextgen", "digest": DIGEST, } ] ), DEV_REGISTRY, ) self.assertEqual(images[0]["repository"], "tidbcloud/pd") self.assertEqual(images[0]["tag"], "v26.3.4-nextgen") self.assertEqual(images[0]["digest"], DIGEST) def test_rejects_source_outside_the_environment_central_ecr(self): payload = json.dumps( [ { "source": ( "380838443567.dkr.ecr.us-west-2.amazonaws.com/" "tidbcloud/pd:v26.3.4-nextgen" ), "digest": DIGEST, } ] ) with self.assertRaisesRegex( MODULE.WorkflowInputError, "configured central ECR" ): MODULE.parse_images(payload, DEV_REGISTRY) def test_rejects_duplicate_sources(self): image = { "source": f"{DEV_REGISTRY}/tidbcloud/pd:v26.3.4-nextgen", "digest": DIGEST, } with self.assertRaisesRegex(MODULE.WorkflowInputError, "duplicated"): MODULE.parse_images(json.dumps([image, image]), DEV_REGISTRY) def test_rejects_invalid_digest(self): payload = json.dumps( [ { "source": f"{DEV_REGISTRY}/tidbcloud/pd:v26.3.4-nextgen", "digest": "sha256:not-a-digest", } ] ) with self.assertRaisesRegex(MODULE.WorkflowInputError, "sha256 digest"): MODULE.parse_images(payload, DEV_REGISTRY) def test_validates_target_registry_and_role_account(self): targets = MODULE.parse_targets_data( [ { "name": "123456789012.dkr.ecr.ap-southeast-1.amazonaws.com", "roleArn": "arn:aws:iam::123456789012:role/byoc-image-sync", } ] ) self.assertEqual(targets[0]["region"], "ap-southeast-1") def test_rejects_target_role_from_another_account(self): with self.assertRaisesRegex(MODULE.WorkflowInputError, "do not match"): MODULE.parse_targets_data( [ { "name": ( "123456789012.dkr.ecr.ap-southeast-1.amazonaws.com" ), "roleArn": ( "arn:aws:iam::210987654321:role/byoc-image-sync" ), } ] ) def test_writes_compact_matrix_outputs(self): with tempfile.NamedTemporaryFile(mode="r+", encoding="utf-8") as output: MODULE.write_github_outputs( output.name, { "images": [{"source": "example", "digest": DIGEST}], "image_count": 1, }, ) output.seek(0) lines = output.read().splitlines() self.assertEqual(lines[0], f'images=[{{"digest":"{DIGEST}","source":"example"}}]') self.assertEqual(lines[1], "image_count=1") def test_requires_a_canonical_uuid_correlation_id(self): MODULE.validate_correlation_id("00000000-0000-4000-8000-000000000000") with self.assertRaisesRegex(MODULE.WorkflowInputError, "UUID"): MODULE.validate_correlation_id("not-a-uuid") def test_aws_role_uses_runner_web_identity_instead_of_static_credentials(self): with tempfile.NamedTemporaryFile() as token: environment = { "AWS_WEB_IDENTITY_TOKEN_FILE": token.name, "AWS_ACCESS_KEY_ID": "must-not-be-used", "AWS_SECRET_ACCESS_KEY": "must-not-be-used", "AWS_SESSION_TOKEN": "must-not-be-used", "GITHUB_RUN_ID": "1234", } with mock.patch.dict(MODULE.os.environ, environment, clear=True): result = MODULE.aws_environment( "arn:aws:iam::123456789012:role/byoc-image-sync", "ap-southeast-1", ) self.assertNotIn("AWS_ACCESS_KEY_ID", result) self.assertNotIn("AWS_SECRET_ACCESS_KEY", result) self.assertNotIn("AWS_SESSION_TOKEN", result) self.assertEqual( result["AWS_ROLE_ARN"], "arn:aws:iam::123456789012:role/byoc-image-sync", ) self.assertEqual(result["AWS_ROLE_SESSION_NAME"], "byoc-sync-1234") def test_requires_the_runner_web_identity_token_file(self): with mock.patch.dict(MODULE.os.environ, {}, clear=True): with self.assertRaisesRegex(RuntimeError, "token file is unavailable"): MODULE.aws_environment( "arn:aws:iam::123456789012:role/byoc-image-sync", "ap-southeast-1", ) def test_refuses_to_overwrite_a_conflicting_target_tag(self): target = { "name": "123456789012.dkr.ecr.ap-southeast-1.amazonaws.com", "roleArn": "arn:aws:iam::123456789012:role/byoc-image-sync", } image = { "source": f"{DEV_REGISTRY}/tidbcloud/pd:v26.3.4-nextgen", "repository": "tidbcloud/pd", "tag": "v26.3.4-nextgen", "digest": DIGEST, } environment = { "TARGET_JSON": json.dumps(target), "SOURCE_ROLE_ARN": "arn:aws:iam::385595570414:role/source", "TARGET_ROLE_ARN": target["roleArn"], } with ( mock.patch.dict(MODULE.os.environ, environment, clear=True), mock.patch.object( MODULE, "normalized_images_from_environment", return_value=[image] ), mock.patch.object(MODULE, "login_ecr", side_effect=[{}, {}]), mock.patch.object(MODULE, "ensure_repository"), mock.patch.object( MODULE, "target_digest", return_value="sha256:" + "b" * 64 ), mock.patch.object(MODULE, "run_checked") as copy, ): with self.assertRaisesRegex(RuntimeError, "refusing to overwrite"): MODULE.command_sync_target( SimpleNamespace(source_registry=DEV_REGISTRY) ) copy.assert_not_called() def test_target_digest_uses_registry_api_and_treats_not_found_as_absent(self): image = ( "123456789012.dkr.ecr.ap-southeast-1.amazonaws.com/" "tidbcloud/pd:v26.3.4-nextgen" ) not_found = SimpleNamespace( returncode=1, stdout="", stderr="failed to get manifest: 404 Not Found", ) with mock.patch.object(MODULE.subprocess, "run", return_value=not_found) as run: self.assertIsNone(MODULE.target_digest(image)) run.assert_called_once_with( ["regctl", "image", "digest", image], check=False, text=True, capture_output=True, ) def test_target_digest_rejects_registry_errors(self): denied = SimpleNamespace( returncode=1, stdout="", stderr="unauthorized: access denied", ) with mock.patch.object(MODULE.subprocess, "run", return_value=denied): with self.assertRaisesRegex(RuntimeError, "access denied"): MODULE.target_digest("registry.example/tidbcloud/pd:v1") def test_copies_from_digest_to_the_same_repository_path_and_tag(self): target = { "name": "123456789012.dkr.ecr.ap-southeast-1.amazonaws.com", "roleArn": "arn:aws:iam::123456789012:role/byoc-image-sync", } image = { "source": f"{DEV_REGISTRY}/tidbcloud/pd:v26.3.4-nextgen", "repository": "tidbcloud/pd", "tag": "v26.3.4-nextgen", "digest": DIGEST, } environment = { "TARGET_JSON": json.dumps(target), "SOURCE_ROLE_ARN": "arn:aws:iam::385595570414:role/source", "TARGET_ROLE_ARN": target["roleArn"], } with ( mock.patch.dict(MODULE.os.environ, environment, clear=True), mock.patch.object( MODULE, "normalized_images_from_environment", return_value=[image] ), mock.patch.object(MODULE, "login_ecr", side_effect=[{}, {}]), mock.patch.object(MODULE, "ensure_repository"), mock.patch.object(MODULE, "target_digest", side_effect=[None, DIGEST]), mock.patch.object(MODULE, "run_checked") as copy, ): MODULE.command_sync_target(SimpleNamespace(source_registry=DEV_REGISTRY)) copy.assert_called_once_with( [ "regctl", "image", "copy", f"{DEV_REGISTRY}/tidbcloud/pd@{DIGEST}", ( "123456789012.dkr.ecr.ap-southeast-1.amazonaws.com/" "tidbcloud/pd:v26.3.4-nextgen" ), ] ) if __name__ == "__main__": unittest.main()