EncrypterParv commented on code in PR #75377:
URL: https://github.com/apache/airflow/pull/75377#discussion_r4393118145
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_security.py:
##########
@@ -604,6 +604,122 @@ async def
test_dag_processor_token_resolves_the_team_of_its_bundle(self, session
assert result == "team_a"
+ @pytest.mark.db_test
+ @pytest.mark.asyncio
+ @pytest.mark.usefixtures("async_db_engine")
+ async def
test_callback_execution_token_resolves_team_from_its_bundle(self, session):
+ team = Team(name=f"team_{uuid4().hex}")
+ bundle = DagBundleModel(name=f"bundle_{uuid4().hex}")
+ bundle.teams.append(team)
+ callback = Callback()
+ callback.bundle_name = bundle.name
+ callback.fetch_method = CallbackFetchMethod.IMPORT_PATH
+ session.add_all([bundle, callback])
+ session.commit()
+ token = TIToken(id=callback.id, claims=TIClaims(scope="execution"))
+
+ try:
+ with conf_vars({("core", "multi_team"): "True"}):
+ result = await get_team_name_dep(token=token)
+ finally:
+ session.delete(callback)
+ session.delete(bundle)
+ session.delete(team)
+ session.commit()
+
+ assert result == team.name
+
+ @pytest.mark.db_test
+ @pytest.mark.asyncio
+ @pytest.mark.usefixtures("async_db_engine")
+ async def test_task_execution_token_still_resolves_team_from_its_dag(self,
session, create_task_instance):
Review Comment:
hey, thanks for the review!!
I've addressed all the suggestions, do let me know if u want to make any
further changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
