Repository: apache/airflow
File: providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py (Python)

Code before:
            ("is_authorized_dag", DagDetails(team_name="team-a"), "Dag:team-a#LIST"),
            ("is_authorized_connection", ConnectionDetails(team_name="team-a"), "Connection:team-a#LIST"),
            ("is_authorized_variable", VariableDetails(team_name="team-a"), "Variable:team-a#LIST"),
            ("is_authorized_pool", PoolDetails(team_name="team-a"), "Pool:team-a#LIST"),
        ]
        if AIRFLOW_V_3_2_PLUS
        else [],
    )
    def test_list_with_team_name_uses_team_scoped_permission(
        self, auth_manager_multi_team, user, function, details, permission
    ):
        user.access_token = _build_access_token({"groups": ["team-a"]})
        mock_response = Mock()
        mock_response.status_code = 200
        auth_manager_multi_team.http_session.post = Mock(return_value=mock_response)

        getattr(auth_manager_multi_team, function)(method="GET", user=user, details=details)

        actual_permission = auth_manager_multi_team.http_session.post.call_args.kwargs["data"]["permission"]
        assert actual_permission == permission

    def test_filter_authorized_dag_ids_team_mismatch(self, auth_manager_multi_team, user):
        if not AIRFLOW_V_3_2_PLUS:
            pytest.skip("team_name not supported before Airflow 3.2.0")

        with patch.object(KeycloakAuthManager, "is_authorized_dag", return_value=False) as mock_is_authorized:
            result = auth_manager_multi_team.filter_authorized_dag_ids(
                dag_ids={"dag-a"}, user=user, team_name="team-b"
            )

            mock_is_authorized.assert_called_once()
            assert result == set()

    def test_filter_authorized_dag_ids_team_match(self, auth_manager_multi_team, user):
        if not AIRFLOW_V_3_2_PLUS:
            pytest.skip("team_name not supported before Airflow 3.2.0")

        with patch.object(KeycloakAuthManager, "is_authorized_dag", return_value=True) as mock_is_authorized:
            result = auth_manager_multi_team.filter_authorized_dag_ids(
                dag_ids={"dag-a"}, user=user, team_name="team-a"
            )

            mock_is_authorized.assert_called_once()
            assert result == {"dag-a"}

    def test_filter_authorized_pools_no_team_returns_empty(self, auth_manager_multi_team, user):
        if not AIRFLOW_V_3_2_PLUS:
            pytest.skip("filter_authorized_pools not supported before Airflow 3.2.0")

        with patch.object(
            KeycloakAuthManager, "is_authorized_pool", return_value=False

Review comment (suggestion):
nit: this is a personal preference but I usually prefer using decorator instead of `with`, that makes the test implementation shorter

Code after:
        [
            ("is_authorized_dag", DagDetails, {"team_name": "team-a"}, "Dag:team-a#LIST"),
            (
                "is_authorized_connection",
                ConnectionDetails,
                {"team_name": "team-a"},
                "Connection:team-a#LIST",
            ),
            ("is_authorized_variable", VariableDetails, {"team_name": "team-a"}, "Variable:team-a#LIST"),
            ("is_authorized_pool", PoolDetails, {"team_name": "team-a"}, "Pool:team-a#LIST"),
        ],
    )
    def test_list_with_team_name_uses_team_scoped_permission(
        self, auth_manager_multi_team, user, function, details_cls, details_kwargs, permission
    ):
        details = details_cls(**details_kwargs)
        user.access_token = _build_access_token({"groups": ["team-a"]})
        mock_response = Mock()
        mock_response.status_code = 200
        auth_manager_multi_team.http_session.post = Mock(return_value=mock_response)

        getattr(auth_manager_multi_team, function)(method="GET", user=user, details=details)

        actual_permission = auth_manager_multi_team.http_session.post.call_args.kwargs["data"]["permission"]
        assert actual_permission == permission

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team_name not supported before Airflow 3.2.0")
    @patch.object(KeycloakAuthManager, "is_authorized_dag", return_value=False)
    def test_filter_authorized_dag_ids_team_mismatch(self, mock_is_authorized, auth_manager_multi_team, user):
        result = auth_manager_multi_team.filter_authorized_dag_ids(
            dag_ids={"dag-a"}, user=user, team_name="team-b"
        )

        mock_is_authorized.assert_called_once()
        assert result == set()

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team_name not supported before Airflow 3.2.0")
    @patch.object(KeycloakAuthManager, "is_authorized_dag", return_value=True)
    def test_filter_authorized_dag_ids_team_match(self, mock_is_authorized, auth_manager_multi_team, user):
        result = auth_manager_multi_team.filter_authorized_dag_ids(
            dag_ids={"dag-a"}, user=user, team_name="team-a"
        )

        mock_is_authorized.assert_called_once()
        assert result == {"dag-a"}

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team_name not supported before Airflow 3.2.0")
    @patch.object(KeycloakAuthManager, "is_authorized_pool", return_value=False)
    def test_filter_authorized_pools_no_team_returns_empty(
        self, mock_is_authorized, auth_manager_multi_team, user
    ):