Test get_upstream_artifacts_with_zip_extract_param

beetmoverscript/task.py

@@ -14,8 +14,8 @@
     RESTRICTED_BUCKET_PATHS,
     CHECKSUMS_CUSTOM_FILE_NAMING
 )
+from scriptworker import artifacts as scriptworker_artifacts
 from scriptworker.exceptions import ScriptWorkerTaskException
-from scriptworker.artifacts import get_and_check_single_upstream_artifact_full_path
 
 log = logging.getLogger(__name__)
 
@@ -116,7 +116,7 @@ def get_upstream_artifacts(context, preserve_full_paths=False):
         locale = artifact_dict['locale']
         artifacts[locale] = artifacts.get(locale, {})
         for path in artifact_dict['paths']:
-            abs_path = get_and_check_single_upstream_artifact_full_path(
+            abs_path = scriptworker_artifacts.get_and_check_single_upstream_artifact_full_path(
                 context, artifact_dict['taskId'], path
             )
             if preserve_full_paths:
@@ -127,22 +127,26 @@ def get_upstream_artifacts(context, preserve_full_paths=False):
 
 
 def get_upstream_artifacts_with_zip_extract_param(context):
-    return [
-        (artifact_definition['taskId'], artifact_definition['paths'], artifact_definition.get('zipExtract', False))
-        for artifact_definition in context.task['payload']['upstreamArtifacts']
-    ]
-
-    # TODO not use a dict comprehension as it erases paths
-    # return {
-    #     task_id: [
-    #         'paths': [
-    #             get_and_check_single_upstream_artifact_full_path(context, task_id, path)
-    #             for path in paths
-    #         ],
-    #         'zip_extract': zip_extract,
-    #     ]
-    #     for task_id, paths, zip_extract in task_ids_and_relative_paths
-    # }
+    # XXX A dict comprehension isn't used because upstream_definition would be erased if the same
+    # taskId is present twice in upstreamArtifacts
+    upstream_artifacts_per_task_id = {}
+
+    for artifact_definition in context.task['payload']['upstreamArtifacts']:
+        task_id = artifact_definition['taskId']
+        upstream_definitions = upstream_artifacts_per_task_id.get(task_id, [])
+
+        new_upstream_definition = {
+            'paths': [
+                scriptworker_artifacts.get_and_check_single_upstream_artifact_full_path(context, task_id, path)
+                for path in artifact_definition['paths']
+            ],
+            'zip_extract': artifact_definition.get('zipExtract', False),
+        }
+
+        upstream_definitions.append(new_upstream_definition)
+        upstream_artifacts_per_task_id[task_id] = upstream_definitions
+
+    return upstream_artifacts_per_task_id
 
 
 def get_release_props(context, platform_mapping=STAGE_PLATFORM_MAP):

beetmoverscript/test/test_task.py

@@ -7,8 +7,8 @@
 )
 from beetmoverscript.task import (
     validate_task_schema, add_balrog_manifest_to_artifacts,
-    get_upstream_artifacts, generate_checksums_manifest,
-    get_task_bucket, get_task_action,
+    get_upstream_artifacts, get_upstream_artifacts_with_zip_extract_param,
+    generate_checksums_manifest, get_task_bucket, get_task_action,
     validate_bucket_paths, get_release_props, is_custom_beetmover_checksums_task
 )
 from scriptworker.context import Context
@@ -55,13 +55,71 @@ def test_get_upstream_artifacts(expected, preserve):
     assert sorted(list(artifacts_to_beetmove['en-US'])) == sorted(expected)
 
 
-def test_get_upstream_artifacts_with_zip_extract_param():
+def test_get_upstream_artifacts_with_zip_extract_param(monkeypatch):
     context = Context()
     context.config = get_fake_valid_config()
     context.task = get_fake_valid_task()
     context.properties = context.task['payload']['releaseProperties']
 
-    assert get_upstream_artifacts_with_zip_extract_param(context) == {}
+    context.task['payload']['upstreamArtifacts'] = [{
+        'paths': ['a/non/archive', 'another/non/archive'],
+        'taskId': 'firstTaskId',
+        'taskType': 'someType1',
+        'zipExtract': False,
+    }, {
+        'paths': ['archive1.zip', 'subfolder/archive2.zip'],
+        'taskId': 'firstTaskId',
+        'taskType': 'someType1',
+        'zipExtract': True,
+    }, {
+        'paths': ['just/another/regular/file'],
+        'taskId': 'secondTaskId',
+        'taskType': 'someType2',
+        'zipExtract': False,
+    }, {
+        'paths': ['archive1.zip'],
+        'taskId': 'thirdTaskId',
+        'taskType': 'someType3',
+        'zipExtract': True,
+    }]
+
+    def mock_upstream_artifact_full_path(context, task_id, path):
+        # doesn't check whether the file exists on disk
+        return os.path.join(context.config['work_dir'], 'cot', task_id, path)
+
+    monkeypatch.setattr(
+        'scriptworker.artifacts.get_and_check_single_upstream_artifact_full_path',
+        mock_upstream_artifact_full_path
+    )
+
+    assert get_upstream_artifacts_with_zip_extract_param(context) == {
+        'firstTaskId': [{
+            'paths': [
+                os.path.join(context.config['work_dir'], 'cot', 'firstTaskId', 'a/non/archive'),
+                os.path.join(context.config['work_dir'], 'cot', 'firstTaskId', 'another/non/archive'),
+            ],
+            'zip_extract': False,
+        }, {
+            'paths': [
+                os.path.join(context.config['work_dir'], 'cot', 'firstTaskId', 'archive1.zip'),
+                os.path.join(context.config['work_dir'], 'cot', 'firstTaskId', 'subfolder/archive2.zip'),
+            ],
+            'zip_extract': True,
+        }],
+        'secondTaskId': [{
+            'paths': [
+                os.path.join(context.config['work_dir'], 'cot', 'secondTaskId', 'just/another/regular/file'),
+            ],
+            'zip_extract': False,
+        }],
+        'thirdTaskId': [{
+            'paths': [
+                os.path.join(context.config['work_dir'], 'cot', 'thirdTaskId', 'archive1.zip'),
+            ],
+            'zip_extract': True,
+        }],
+    }
+
 
 # validate_task {{{1
 def test_validate_task(context):

tox.ini

@@ -25,7 +25,7 @@ deps =
     virtualenv
 
 commands=
-    py.test --cov-config .coveragerc --cov=beetmoverscript --cov-report term-missing -- beetmoverscript/test/test_task.py
+    py.test --cov-config .coveragerc --cov=beetmoverscript --cov-report term-missing
     coverage html
     flake8