def test_session_header_decode_success(self, instance, coordinator,
workspace, external_pipeline):
run_id, jwt_header, expected_email = (
"foo",
"foo.eyJlbWFpbCI6ICJoZWxsb0BlbGVtZW50bC5jb20ifQ==.bar",
"*****@*****.**",
)
MockRequest = namedtuple("MockRequest", ["headers"])
workspace._source = MockRequest(headers={
"X-Amzn-Trace-Id": "some_info",
"X-Amzn-Oidc-Data": jwt_header,
})
run = self.create_run(
instance,
external_pipeline,
run_id=run_id,
status=PipelineRunStatus.NOT_STARTED,
)
returned_run = coordinator.submit_run(SubmitRunContext(run, workspace))
assert returned_run.run_id == run_id
assert returned_run.status == PipelineRunStatus.QUEUED
tags = instance.get_run_tags()
assert len(tags) == 1
(tag_name, set_of_tag_values) = tags[0]
assert tag_name == "user"
assert set_of_tag_values == {expected_email}
def test_session_header_decode_success(self, instance, coordinator,
workspace, external_pipeline):
run_id, jwt_header, expected_email = (
"foo",
"foo.eyJlbWFpbCI6ICJoZWxsb0BlbGVtZW50bC5jb20ifQ==.bar",
"*****@*****.**",
)
with patch(
"run_attribution_example.custom_run_coordinator.has_request_context"
) as mock_has_request_context, patch(
"run_attribution_example.custom_run_coordinator.request",
headers={
"X-Amzn-Trace-Id": "some_info",
"X-Amzn-Oidc-Data": jwt_header,
},
):
mock_has_request_context.return_value = True
run = self.create_run(
instance,
external_pipeline,
run_id=run_id,
status=PipelineRunStatus.NOT_STARTED,
)
returned_run = coordinator.submit_run(
SubmitRunContext(run, workspace))
assert returned_run.run_id == run_id
assert returned_run.status == PipelineRunStatus.QUEUED
tags = instance.get_run_tags()
assert len(tags) == 1
(tag_name, set_of_tag_values) = tags[0]
assert tag_name == "user"
assert set_of_tag_values == {expected_email}
def test_session_header_decode_failure(self, instance, coordinator,
workspace, external_pipeline):
run_id = "foo-1"
with patch(
"run_attribution_example.custom_run_coordinator.has_request_context"
) as mock_has_request_context, patch(
"run_attribution_example.custom_run_coordinator.warnings"
) as mock_warnings:
mock_has_request_context.return_value = False
run = self.create_run(
instance,
external_pipeline,
run_id=run_id,
status=PipelineRunStatus.NOT_STARTED,
)
returned_run = coordinator.submit_run(
SubmitRunContext(run, workspace))
assert returned_run.run_id == run_id
assert returned_run.status == PipelineRunStatus.QUEUED
tags = instance.get_run_tags()
assert len(tags) == 0
mock_warnings.warn.assert_called_once()
assert mock_warnings.warn.call_args.args[0].startswith(
"Couldn't decode JWT header")
def test_submit_run_checks_status(
self, instance, coordinator, workspace, external_pipeline
): # pylint: disable=redefined-outer-name
run = self.create_run(
instance, external_pipeline, run_id="foo-1", status=PipelineRunStatus.QUEUED
)
with pytest.raises(CheckError):
coordinator.submit_run(SubmitRunContext(run, workspace))
def submit_run(self, context: SubmitRunContext) -> PipelineRun:
pipeline_run = context.pipeline_run
jwt_claims_header = context.get_request_header("X-Amzn-Oidc-Data")
email = self.get_email(jwt_claims_header)
if email:
self._instance.add_run_tags(pipeline_run.run_id, {"user": email})
else:
warnings.warn(f"Couldn't decode JWT header {jwt_claims_header}")
return super().submit_run(context)
def test_submit_run_checks_status(instance, coodinator): # pylint: disable=redefined-outer-name
with get_bar_workspace(instance) as workspace:
external_pipeline = (
workspace.get_repository_location("bar_repo_location").
get_repository("bar_repo").get_full_external_pipeline("foo"))
run = create_run(instance,
external_pipeline,
run_id="foo-1",
status=PipelineRunStatus.STARTED)
with pytest.raises(CheckError):
coodinator.submit_run(SubmitRunContext(run, workspace))
def test_submit_run(
self, instance, coordinator, workspace, external_pipeline
): # pylint: disable=redefined-outer-name
run = self.create_run(
instance, external_pipeline, run_id="foo-1", status=PipelineRunStatus.NOT_STARTED
)
returned_run = coordinator.submit_run(SubmitRunContext(run, workspace))
assert returned_run.run_id == "foo-1"
assert returned_run.status == PipelineRunStatus.QUEUED
assert len(instance.run_launcher.queue()) == 0
stored_run = instance.get_run_by_id("foo-1")
assert stored_run.status == PipelineRunStatus.QUEUED
def test_cancel_run(
self, instance, coordinator, workspace, external_pipeline
): # pylint: disable=redefined-outer-name
run = self.create_run(
instance, external_pipeline, run_id="foo-1", status=PipelineRunStatus.NOT_STARTED
)
assert not coordinator.can_cancel_run(run.run_id)
coordinator.submit_run(SubmitRunContext(run, workspace))
assert coordinator.can_cancel_run(run.run_id)
coordinator.cancel_run(run.run_id)
stored_run = instance.get_run_by_id("foo-1")
assert stored_run.status == PipelineRunStatus.CANCELED
assert not coordinator.can_cancel_run(run.run_id)
def test_submit_run(instance, coodinator): # pylint: disable=redefined-outer-name
with get_bar_workspace(instance) as workspace:
external_pipeline = (
workspace.get_repository_location("bar_repo_location").
get_repository("bar_repo").get_full_external_pipeline("foo"))
run = create_run(instance, external_pipeline, run_id="foo-1")
returned_run = coodinator.submit_run(SubmitRunContext(run, workspace))
assert returned_run.run_id == "foo-1"
assert returned_run.status == PipelineRunStatus.STARTING
assert len(instance.run_launcher.queue()) == 1
assert instance.run_launcher.queue()[0].run_id == "foo-1"
assert instance.run_launcher.queue(
)[0].status == PipelineRunStatus.STARTING
assert instance.get_run_by_id("foo-1")
def submit_run(self, context: SubmitRunContext) -> PipelineRun:
desired_header = context.get_request_header(CUSTOM_HEADER_NAME)