Eason09053360 commented on code in PR #74364:
URL: https://github.com/apache/airflow/pull/74364#discussion_r4213843339
##########
airflow-core/tests/unit/models/test_revoked_token.py:
##########
@@ -97,7 +104,117 @@ def test_cleanup_skips_when_interval_not_passed(self):
):
RevokedToken.is_revoked("test-jti", session=mock_session)
- # session.execute should NOT be called
+ mock_session.scalars.assert_not_called()
+ mock_session.execute.assert_not_called()
+ finally:
+ RevokedToken._last_cleanup_time = original_last_cleanup
+
+ def test_cleanup_skipped_while_another_thread_is_cleaning(self):
+ """The interval bookkeeping is not thread safe, so only one pass may
run at a time."""
+ mock_session = MagicMock()
+ mock_session.scalar.return_value = False
+
+ original_last_cleanup = RevokedToken._last_cleanup_time
+ RevokedToken._cleanup_lock.acquire()
+ try:
+ RevokedToken._last_cleanup_time = 0.0
+ with (
+ patch("airflow.models.revoked_token.time.monotonic",
return_value=8000.0),
+ patch("airflow.models.revoked_token.conf.getint",
return_value=3600),
+ ):
+ assert RevokedToken.is_revoked("test-jti",
session=mock_session) is False
+
+ mock_session.scalars.assert_not_called()
mock_session.execute.assert_not_called()
+ # a skipped pass must not claim the interval either
+ assert RevokedToken._last_cleanup_time == 0.0
finally:
+ RevokedToken._cleanup_lock.release()
RevokedToken._last_cleanup_time = original_last_cleanup
+
+ def
test_failed_cleanup_rolls_back_so_the_revocation_read_still_works(self):
+ """A failed statement aborts the transaction on PostgreSQL; the read
after it must not inherit that."""
+ mock_session = MagicMock()
+ mock_session.scalar.return_value = False
+ mock_session.scalars.side_effect = RuntimeError("database is on fire")
+
+ original_last_cleanup = RevokedToken._last_cleanup_time
+ try:
+ RevokedToken._last_cleanup_time = 0.0
+ with (
+ patch("airflow.models.revoked_token.time.monotonic",
return_value=8000.0),
+ patch("airflow.models.revoked_token.conf.getint",
return_value=3600),
+ ):
+ assert RevokedToken.is_revoked("test-jti",
session=mock_session) is False
+
+ mock_session.rollback.assert_called_once()
+ # the lock must not stay held after a failure
+ assert RevokedToken._cleanup_lock.acquire(blocking=False)
+ RevokedToken._cleanup_lock.release()
+ finally:
+ RevokedToken._last_cleanup_time = original_last_cleanup
+
+
[email protected]_test
+class TestRevokedTokenCleanupIsBounded:
+ """Cleanup runs on the request path, so a single pass must not issue an
unbounded DELETE."""
+
+ @pytest.fixture(autouse=True)
+ def reset_cleanup_state(self):
+ original_last_cleanup = RevokedToken._last_cleanup_time
+ with create_session() as session:
+ session.execute(delete(RevokedToken))
+ yield
+ RevokedToken._last_cleanup_time = original_last_cleanup
+ with create_session() as session:
+ session.execute(delete(RevokedToken))
+
+ @staticmethod
+ def _add_tokens(expired: int, live: int) -> None:
+ now = datetime.now(tz=timezone.utc)
+ with create_session() as session:
+ for i in range(expired):
+ session.add(RevokedToken(jti=f"expired-{i}", exp=now -
timedelta(hours=1)))
+ for i in range(live):
+ session.add(RevokedToken(jti=f"live-{i}", exp=now +
timedelta(hours=1)))
+
+ @staticmethod
+ def _remaining() -> int:
+ with create_session() as session:
+ return
session.scalars(select(func.count()).select_from(RevokedToken)).one()
+
+ @conf_vars({("api_auth", "jwt_expiration_time"): "3600"})
+ def test_cleanup_deletes_at_most_one_batch(self):
+ self._add_tokens(expired=7, live=2)
+ RevokedToken._last_cleanup_time = 0.0
+
+ with (
+ patch("airflow.models.revoked_token._CLEANUP_BATCH_SIZE", 3),
+ patch("airflow.models.revoked_token.time.monotonic",
return_value=100_000.0),
+ ):
+ RevokedToken.is_revoked("live-0")
+
+ # 3 of the 7 expired rows gone, both unexpired rows untouched
+ assert self._remaining() == 6
+
+ @conf_vars({("api_auth", "jwt_expiration_time"): "3600"})
+ def test_full_batch_lets_the_next_check_resume_draining(self):
+ self._add_tokens(expired=7, live=0)
+ RevokedToken._last_cleanup_time = 0.0
+
+ with (
+ patch("airflow.models.revoked_token._CLEANUP_BATCH_SIZE", 3),
+ patch("airflow.models.revoked_token.time.monotonic",
return_value=100_000.0),
+ ):
+ # Each full batch rewinds the interval, so the passes chain on a
frozen clock.
+ RevokedToken.is_revoked("expired-0")
+ assert self._remaining() == 4
+ RevokedToken.is_revoked("expired-0")
+ assert self._remaining() == 1
+ RevokedToken.is_revoked("expired-0")
+ assert self._remaining() == 0
+
+ # The last pass did not fill the batch, so the interval applies
again
+ self._add_tokens(expired=2, live=0)
+ RevokedToken.is_revoked("expired-0")
+ assert self._remaining() == 2
Review Comment:
Added
--
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]
