From 6508f8b8802a113e15f9b426f6c558b994c9eb1a Mon Sep 17 00:00:00 2001
From: Ben Hearsum <ben@mozilla.com>
Date: Mon, 2 Jun 2025 13:58:25 -0400
Subject: [PATCH] feat: add payload builder for beetmover upload-data action

This is a companion to https://github.com/mozilla-releng/scriptworker-scripts/pull/1222 where this beetmover feature was added.
---
 src/mozilla_taskgraph/worker_types.py | 54 ++++++++++++++++++++++
 test/test_worker_types.py             | 66 +++++++++++++++++++++++++++
 2 files changed, 120 insertions(+)

diff --git a/src/mozilla_taskgraph/worker_types.py b/src/mozilla_taskgraph/worker_types.py
index 76f9b12..ce37271 100644
--- a/src/mozilla_taskgraph/worker_types.py
+++ b/src/mozilla_taskgraph/worker_types.py
@@ -1,3 +1,6 @@
+import binascii
+from base64 import b64decode
+
 from taskgraph.transforms.task import payload_builder, taskref_or_string
 from voluptuous import Any, Extra, Optional, Required
 
@@ -463,3 +466,54 @@ def dash_to_underscore(obj):
     for k, v in obj.items():
         new_obj[k.replace("-", "_")] = v
     return new_obj
+
+
+@payload_builder(
+    "scriptworker-beetmover-data",
+    schema={
+        Required("app-name"): str,
+        Required("bucket"): str,
+        Required("project"): str,
+        Required("data-map"): [
+            {
+                Required("data"): str,
+                Required("content-type"): str,
+                Required("destinations"): [str],
+            },
+        ],
+        Optional("dry-run"): bool,
+    },
+)
+def build_beetmover_data_payload(_, task, task_def):
+    worker = task["worker"]
+    task_def["tags"]["worker-implementation"] = "scriptworker"
+
+    task_def["payload"] = {
+        "releaseProperties": {
+            "appName": worker["app-name"],
+        },
+        "dataMap": [],
+    }
+
+    for dataMap in worker["data-map"]:
+        data = dataMap["data"]
+        try:
+            b64decode(data)
+        except binascii.Error:
+            raise Exception(f"data must be base64 encoded! data was: {data}")
+
+        task_def["payload"]["dataMap"].append(
+            {
+                "data": data,
+                "contentType": dataMap["content-type"],
+                "destinations": dataMap["destinations"],
+            }
+        )
+
+    scopes = task_def.setdefault("scopes", [])
+    scopes.extend(
+        [
+            f"project:{worker['project']}:releng:beetmover:bucket:{worker['bucket']}",
+            f"project:{worker['project']}:releng:beetmover:action:upload-data",
+        ]
+    )
diff --git a/test/test_worker_types.py b/test/test_worker_types.py
index ae8eba7..96912fe 100644
--- a/test/test_worker_types.py
+++ b/test/test_worker_types.py
@@ -1115,3 +1115,69 @@ def test_lando_merge_beta_to_release(build_payload):
         ],
         "tags": {"worker-implementation": "scriptworker"},
     }
+
+
+def test_beetmover_upload_data(build_payload):
+    worker = {
+        "bucket": "testbucket",
+        "app-name": "testapp",
+        "project": "testproject",
+        "data-map": [
+            {
+                "data": "c29tZSB0ZXN0IGRhdGE=",
+                "content-type": "text/plain",
+                "destinations": [
+                    "foo.txt",
+                    "bar.txt",
+                ],
+            },
+        ],
+    }
+    _, task_def = build_payload("scriptworker-beetmover-data", worker=worker)
+    assert task_def == {
+        "payload": {
+            "releaseProperties": {
+                "appName": "testapp",
+            },
+            "dataMap": [
+                {
+                    "data": "c29tZSB0ZXN0IGRhdGE=",
+                    "contentType": "text/plain",
+                    "destinations": [
+                        "foo.txt",
+                        "bar.txt",
+                    ],
+                },
+            ],
+        },
+        "scopes": [
+            "project:testproject:releng:beetmover:bucket:testbucket",
+            "project:testproject:releng:beetmover:action:upload-data",
+        ],
+        "tags": {
+            "worker-implementation": "scriptworker",
+        },
+    }
+
+
+def test_beetmover_upload_data_bad_encoding(build_payload):
+    worker = {
+        "bucket": "testbucket",
+        "app-name": "testapp",
+        "project": "testproject",
+        "data-map": [
+            {
+                "data": "aaaaa",
+                "content-type": "text/plain",
+                "destinations": [
+                    "foo.txt",
+                    "bar.txt",
+                ],
+            },
+        ],
+    }
+    try:
+        build_payload("scriptworker-beetmover-data", worker=worker)
+        assert False, "should've raised an exception"
+    except Exception as e:
+        assert "data must be base64 encoded" in e.args[0]