Skip to content

Commit

Permalink
use context.temp_queue to download artifacts
Browse files Browse the repository at this point in the history
This will break if `context.temp_queue` is None.
  • Loading branch information
escapewindow committed May 3, 2022
1 parent da8e6df commit b8b0c2b
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 13 deletions.
4 changes: 2 additions & 2 deletions src/scriptworker/artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,9 +207,9 @@ def get_artifact_url(context, task_id, path):
"""
if path.startswith("public/"):
url = context.queue.buildUrl("getLatestArtifact", task_id, path)
url = context.temp_queue.buildUrl("getLatestArtifact", task_id, path)
else:
url = context.queue.buildSignedUrl(
url = context.temp_queue.buildSignedUrl(
"getLatestArtifact",
task_id,
path,
Expand Down
25 changes: 21 additions & 4 deletions src/scriptworker/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,6 @@ class Context(object):
containing the scriptworker credentials.
session (aiohttp.ClientSession): the default aiohttp session
task (dict): the task definition for the current task.
temp_queue (taskcluster.aio.Queue): the taskcluster Queue object
containing the task-specific temporary credentials.
"""

Expand All @@ -60,8 +58,8 @@ class Context(object):
queue: Optional[Queue] = None
session: Optional[aiohttp.ClientSession] = None
task: Optional[Dict[str, Any]] = None
temp_queue: Optional[Queue] = None
running_tasks = None
_temp_queue = None
_download_semaphore = None
_credentials: Optional[Dict[str, Any]] = None
_claim_task: Optional[Dict[str, Any]] = None # This assumes a single task per worker.
Expand Down Expand Up @@ -185,7 +183,10 @@ def temp_credentials(self) -> Optional[Dict[str, Any]]:
@temp_credentials.setter
def temp_credentials(self, credentials: Optional[Dict[str, Any]]) -> None:
self._temp_credentials = credentials
self.temp_queue = self.create_queue(self.temp_credentials)
if credentials is not None:
self.temp_queue = self.create_queue(self.temp_credentials)
else:
self.temp_queue = None

def write_json(self, path: str, contents: Dict[str, Any], message: str) -> None:
"""Write json to disk.
Expand Down Expand Up @@ -232,6 +233,22 @@ def event_loop(self) -> asyncio.AbstractEventLoop:
def event_loop(self, event_loop: asyncio.AbstractEventLoop) -> None:
self._event_loop = event_loop

@property
def temp_queue(self):
"""dict: The queue for the current task.
These will have different sets of scopes than the worker queue.
"""
if self._temp_queue:
return self._temp_queue
else:
return self.queue

@temp_queue.setter
def temp_queue(self, queue):
self._temp_queue = queue

async def populate_projects(self, force: bool = False) -> None:
"""Download the ``projects.yml`` file and populate ``self.projects``.
Expand Down
8 changes: 4 additions & 4 deletions tests/test_artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,10 +183,10 @@ def buildSignedUrl(*args, **kwargs):
return expected

context = mock.MagicMock()
context.queue = mock.MagicMock()
context.queue.options = {"baseUrl": "https://netloc/"}
context.queue.buildUrl = buildUrl
context.queue.buildSignedUrl = buildSignedUrl
context.temp_queue = mock.MagicMock()
context.temp_queue.options = {"baseUrl": "https://netloc/"}
context.temp_queue.buildUrl = buildUrl
context.temp_queue.buildSignedUrl = buildSignedUrl
assert get_artifact_url(context, "x", path) == expected


Expand Down
6 changes: 6 additions & 0 deletions tests/test_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,12 @@ def test_temp_queue(rw_context, mocker):
assert taskcluster.aio.Queue.called_once_with(
options={"rootUrl": rw_context.config["taskcluster_root_url"], "credentials": rw_context.temp_credentials}, session=rw_context.session
)
assert rw_context._temp_queue is not None
rw_context.temp_queue = None
fake_queue = mocker.MagicMock()
rw_context.queue = fake_queue
assert rw_context._temp_queue is None
assert rw_context.temp_queue is fake_queue


@pytest.mark.asyncio
Expand Down
6 changes: 3 additions & 3 deletions tests/test_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,8 +145,8 @@ async def get_temp_creds_context(config_override=None):
yield context


async def create_task(context, task_id, task_group_id):
payload = integration_create_task_payload(context.config, task_group_id)
async def create_task(context, task_id, task_group_id, **kwargs):
payload = integration_create_task_payload(context.config, task_group_id, **kwargs)
return await context.queue.createTask(task_id, payload)


Expand Down Expand Up @@ -328,7 +328,7 @@ async def test_private_artifacts(context_function):
task_group_id = task_id = slugid.nice()
override = {"task_script": ("bash", "-c", ">&2 echo")}
async with context_function(override) as context:
result = await create_task(context, task_id, task_group_id)
result = await create_task(context, task_id, task_group_id, scopes=["queue:get-artifact:SampleArtifacts/_/X.txt"])
assert result["status"]["state"] == "pending"
path = os.path.join(context.config["artifact_dir"], "SampleArtifacts/_/X.txt")
utils.makedirs(os.path.dirname(path))
Expand Down

0 comments on commit b8b0c2b

Please sign in to comment.