robertpofuk commented on code in PR #72262:
URL: https://github.com/apache/airflow/pull/72262#discussion_r4016199013


##########
providers/edge3/tests/unit/edge3/worker_api/test_auth.py:
##########
@@ -150,3 +155,220 @@ def 
test_strips_edge_worker_v1_prefix_and_falls_back_to_full_path(
         jwt_token_authorization_rest(request, authorization="some-token")
 
         mock_jwt_token_authorization.assert_called_once_with(expected_method, 
"some-token")
+
+
+class TestOidcConfig:
+    """Config readers translate raw ``[edge]`` options into validator 
inputs."""
+
+    @pytest.mark.parametrize(
+        ("configured", "expected"),
+        [
+            pytest.param(None, ["RS256"], id="default-rs256"),
+            pytest.param("RS512", ["RS512"], id="single-algorithm"),
+            pytest.param("RS256, RS512", ["RS256", "RS512"], 
id="comma-list-is-split-and-stripped"),
+        ],
+    )
+    def test_oidc_algorithms_parsing(self, configured, expected):
+        """Unset falls back to RS256; a configured value is split on commas 
and stripped."""
+        overrides = {} if configured is None else {("edge", "jwt_algorithm"): 
configured}
+        with conf_vars(overrides):
+            assert auth._jwt_algorithms() == expected
+
+    @pytest.mark.parametrize(
+        ("configured", "expected"),
+        [
+            pytest.param(None, None, 
id="empty-is-none-accepts-only-aud-less-tokens"),
+            pytest.param("api", "api", id="configured-value-passed-through"),
+        ],
+    )
+    def test_oidc_audience_parsing(self, configured, expected):
+        """An empty audience becomes ``None`` (accept only aud-less tokens); a 
set value is forwarded."""
+        overrides = {} if configured is None else {("edge", "jwt_audience"): 
configured}
+        with conf_vars(overrides):
+            assert auth._jwt_audience() == expected
+
+    @pytest.mark.parametrize(
+        ("configured", "expected"),
+        [
+            pytest.param(None, None, id="empty-is-none-skips-issuer-check"),
+            pytest.param(
+                "https://idp.example.com";, "https://idp.example.com";, 
id="configured-passed-through"
+            ),
+        ],
+    )
+    def test_oidc_issuer_parsing(self, configured, expected):
+        """An empty issuer becomes ``None`` (skip issuer check); a set value 
is forwarded verbatim."""
+        overrides = {} if configured is None else {("edge", "jwt_issuer"): 
configured}
+        with conf_vars(overrides):
+            assert auth._jwt_issuer() == expected
+
+    @pytest.mark.parametrize(
+        ("configured", "expected"),
+        [
+            pytest.param(None, 30, id="default-30"),
+            pytest.param("90", 90, id="configured-value-parsed-as-int"),
+        ],
+    )
+    def test_oidc_leeway_parsing(self, configured, expected):
+        """Unset falls back to 30 seconds; a configured value is read as an 
integer."""
+        overrides = {} if configured is None else {("edge", "jwt_leeway"): 
configured}
+        with conf_vars(overrides):
+            assert auth._jwt_leeway() == expected
+
+
+class TestJwtValidatorSelection:
+    """``jwt_validator`` picks shared-secret vs OIDC based on 
``trusted_jwks_url``."""
+
+    @conf_vars({("api_auth", "jwt_secret"): "secret"})
+    def test_uses_shared_secret_validator_when_oidc_jwks_url_unset(self):
+        """Default path is unchanged: a shared-secret validator with no JWKS 
is built."""
+        validator = auth.jwt_validator()
+
+        assert validator.secret_key == "secret"
+
+        assert validator.jwks is None
+
+    @conf_vars(
+        {
+            ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+            ("edge", "jwt_issuer"): "https://idp.example.com";,
+            ("edge", "jwt_algorithm"): "RS512",
+            ("edge", "jwt_leeway"): "90",
+        }
+    )
+    def test_uses_oidc_validator_when_jwks_url_set(self):
+        """Setting ``trusted_jwks_url`` builds a JWKS-backed validator wired 
from ``[edge]`` config."""
+        validator = auth.jwt_validator()
+
+        assert validator.jwks is not None
+
+        assert validator.algorithm == ["RS512"]
+
+        assert validator.issuer == "https://idp.example.com";
+
+        assert validator.audience is None
+
+        assert validator.leeway == 90
+
+    @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL})
+    def test_empty_issuer_without_verifier_fails_closed(self):
+        """Skipping issuer verification without a ``jwt_verifier`` is a hard 
config error."""
+        with pytest.raises(AirflowConfigException, match="jwt_verifier"):
+            auth.jwt_validator()
+
+    @conf_vars(
+        {
+            ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+            ("edge", "jwt_verifier"): 
"my_company.edge_auth.verify_worker_token",
+        }
+    )
+    def test_empty_issuer_with_verifier_is_allowed(self):
+        """A configured ``jwt_verifier`` permits skipping issuer 
verification."""
+        validator = auth.jwt_validator()
+
+        assert validator.jwks is not None
+
+        assert validator.issuer is None
+
+
+class TestMethodClaimCheck:
+    """``_check_method_claim`` enforces the signed ``method`` only for 
shared-secret tokens."""
+
+    @conf_vars({("api_auth", "jwt_secret"): "secret"})
+    def test_shared_secret_rejects_mismatched_method(self):
+        """A shared-secret token minted for another endpoint is forbidden 
(403)."""
+        with pytest.raises(HTTPException) as exc_info:
+            auth._check_method_claim("worker/register", {"method": 
"worker/other"})
+
+        assert exc_info.value.status_code == 403
+
+    @conf_vars({("api_auth", "jwt_secret"): "secret"})
+    def test_shared_secret_accepts_matching_method(self):
+        """A shared-secret token whose ``method`` matches the request passes 
without raising."""
+        auth._check_method_claim("worker/register", {"method": 
"worker/register"})
+
+    @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL})
+    def test_oidc_skips_method_claim(self):
+        """OIDC tokens carry no ``method`` claim, so the check is skipped 
rather than 403."""
+        auth._check_method_claim("worker/register", {})
+
+
+def _signed_token(audience: str = "") -> str:
+    """Mint a real HS512 token, omitting ``aud`` when passed empty."""
+    generator = JWTGenerator(secret_key=JWT_SECRET, valid_for=300, 
audience=audience)
+    return generator.generate()
+
+
+class TestOidcAudienceVerification:
+    """
+    Pin the real audience semantics against a signed token.
+
+    ``jwt_audience`` empty means ``audience=None``, which PyJWT accepts only 
for
+    tokens that carry no ``aud`` claim and rejects for tokens that do. A 
configured
+    value requires a matching ``aud``.
+    """
+
+    def test_empty_audience_accepts_aud_less_token(self):
+        with conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL}):
+            validator = JWTValidator(secret_key=JWT_SECRET, 
audience=auth._jwt_audience(), leeway=5)

Review Comment:
   Adjusted tests for this case



-- 
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]

Reply via email to