tests/integration_tests/tasks/test_timeout.py - superset - Git at Google

 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
 """Integration tests for GTF timeout handling.

 Uses module-level task functions with manual registry (like test_event_handlers.py)
 to avoid mypy issues with the @task decorator's complex generic types.

 NOTE: Tests that use background threads (timeout/abort handlers) are skipped in
 SQLite environments because SQLite connections cannot be shared across threads.
 """

 from __future__ import annotations

 import time
 import uuid
 from typing import Any

 import pytest
 from superset_core.tasks.types import TaskScope, TaskStatus

 from superset.commands.tasks.cancel import CancelTaskCommand
 from superset.daos.tasks import TaskDAO
 from superset.extensions import db
 from superset.models.tasks import Task
 from superset.tasks.ambient_context import get_context
 from superset.tasks.registry import TaskRegistry
 from superset.tasks.scheduler import execute_task
 from tests.integration_tests.base_tests import SupersetTestCase
 from tests.integration_tests.constants import ADMIN_USERNAME


 def _skip_if_sqlite() -> None:
     """Skip test if running with SQLite database.

     SQLite connections cannot be shared across threads, which breaks
     timeout tests that use background threads for abort handlers.
     Must be called from within a test method (with app context).
     """
     if "sqlite" in db.engine.url.drivername:
         pytest.skip("SQLite connections cannot be shared across threads")


 # Module-level state to track handler calls
 _handler_state: dict[str, Any] = {}


 def _reset_handler_state() -> None:
     """Reset handler state before each test."""
     global _handler_state
     _handler_state = {
         "abort_called": False,
         "handler_exception": None,
     }


 def timeout_abortable_task() -> None:
     """Task with abort handler that exits when aborted."""
     ctx = get_context()

     @ctx.on_abort
     def on_abort() -> None:
         _handler_state["abort_called"] = True

     # Poll for abort signal
     for _ in range(50):
         if _handler_state["abort_called"]:
             return
         time.sleep(0.1)


 def timeout_handler_fails_task() -> None:
     """Task with abort handler that throws an exception."""
     ctx = get_context()

     @ctx.on_abort
     def on_abort() -> None:
         _handler_state["abort_called"] = True
         raise ValueError("Handler crashed!")

     # Sleep longer than timeout
     time.sleep(5)


 def simple_task_with_abort() -> None:
     """Simple task with abort handler for testing."""
     ctx = get_context()

     @ctx.on_abort
     def on_abort() -> None:
         pass


 def quick_task_with_abort() -> None:
     """Quick task that completes before timeout."""
     ctx = get_context()

     @ctx.on_abort
     def on_abort() -> None:
         pass

     time.sleep(0.2)


 def _register_test_tasks() -> None:
     """Register test task functions if not already registered.

     Called in setUp() to ensure tasks are registered regardless of
     whether other tests have cleared the registry.
     """
     registrations = [
         ("test_timeout_abortable", timeout_abortable_task),
         ("test_timeout_handler_fails", timeout_handler_fails_task),
         ("test_timeout_simple", simple_task_with_abort),
         ("test_timeout_quick", quick_task_with_abort),
     ]
     for name, func in registrations:
         if not TaskRegistry.is_registered(name):
             TaskRegistry.register(name, func)


 class TestTimeoutHandling(SupersetTestCase):
     """E2E tests for task timeout functionality."""

     def setUp(self) -> None:
         """Set up test fixtures."""
         super().setUp()
         self.login(ADMIN_USERNAME)
         _register_test_tasks()
         _reset_handler_state()

     def test_timeout_with_abort_handler_results_in_timed_out_status(self) -> None:
         """Task with timeout and abort handler should end with TIMED_OUT status."""
         _skip_if_sqlite()

         # Create task with timeout
         task_obj = TaskDAO.create_task(
             task_type="test_timeout_abortable",
             task_key=f"test_key_{uuid.uuid4().hex[:8]}",
             task_name="Test Timeout",
             scope=TaskScope.SYSTEM,
             properties={"timeout": 1},  # 1 second timeout
         )

         # Execute task via Celery executor (synchronously)
         # Use str(uuid) since Celery serializes args as JSON strings
         result = execute_task.apply(
             args=[str(task_obj.uuid), "test_timeout_abortable", (), {}]
         )

         # Verify execution completed
         assert result.successful()
         assert result.result["status"] == TaskStatus.TIMED_OUT.value

         # Verify abort handler was called
         assert _handler_state["abort_called"]

     def test_user_abort_results_in_aborted_status(self) -> None:
         """User-initiated abort on pending task should result in ABORTED."""
         # Create task (pending state)
         task_obj = TaskDAO.create_task(
             task_type="test_timeout_simple",
             task_key=f"test_key_{uuid.uuid4().hex[:8]}",
             task_name="Test Abort Task",
             scope=TaskScope.SYSTEM,
         )

         # Cancel before execution (pending task abort)
         CancelTaskCommand(task_obj.uuid, force=True).run()

         # Refresh from DB
         db.session.expire_all()
         task_obj = db.session.query(Task).filter_by(uuid=task_obj.uuid).first()
         assert task_obj.status == TaskStatus.ABORTED.value

     def test_no_timeout_when_not_configured(self) -> None:
         """Task without timeout should run to completion regardless of duration."""
         task_obj = TaskDAO.create_task(
             task_type="test_timeout_quick",
             task_key=f"test_key_{uuid.uuid4().hex[:8]}",
             task_name="Test No Timeout",
             scope=TaskScope.SYSTEM,
             # No timeout property
         )

         # Use str(uuid) since Celery serializes args as JSON strings
         result = execute_task.apply(
             args=[str(task_obj.uuid), "test_timeout_quick", (), {}]
         )

         assert result.successful()
         assert result.result["status"] == TaskStatus.SUCCESS.value

     def test_abort_handler_exception_results_in_failure(self) -> None:
         """If abort handler throws during timeout, task should be FAILURE."""
         _skip_if_sqlite()

         task_obj = TaskDAO.create_task(
             task_type="test_timeout_handler_fails",
             task_key=f"test_key_{uuid.uuid4().hex[:8]}",
             task_name="Test Handler Fails",
             scope=TaskScope.SYSTEM,
             properties={"timeout": 1},  # 1 second timeout
         )

         # Use str(uuid) since Celery serializes args as JSON strings
         result = execute_task.apply(
             args=[str(task_obj.uuid), "test_timeout_handler_fails", (), {}]
         )

         assert result.successful()
         assert result.result["status"] == TaskStatus.FAILURE.value
         assert _handler_state["abort_called"]
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.
	"""Integration tests for GTF timeout handling.

	Uses module-level task functions with manual registry (like test_event_handlers.py)
	to avoid mypy issues with the @task decorator's complex generic types.

	NOTE: Tests that use background threads (timeout/abort handlers) are skipped in
	SQLite environments because SQLite connections cannot be shared across threads.
	"""

	from __future__ import annotations

	import time
	import uuid
	from typing import Any

	import pytest
	from superset_core.tasks.types import TaskScope, TaskStatus

	from superset.commands.tasks.cancel import CancelTaskCommand
	from superset.daos.tasks import TaskDAO
	from superset.extensions import db
	from superset.models.tasks import Task
	from superset.tasks.ambient_context import get_context
	from superset.tasks.registry import TaskRegistry
	from superset.tasks.scheduler import execute_task
	from tests.integration_tests.base_tests import SupersetTestCase
	from tests.integration_tests.constants import ADMIN_USERNAME


	def _skip_if_sqlite() -> None:
	"""Skip test if running with SQLite database.

	SQLite connections cannot be shared across threads, which breaks
	timeout tests that use background threads for abort handlers.
	Must be called from within a test method (with app context).
	"""
	if "sqlite" in db.engine.url.drivername:
	pytest.skip("SQLite connections cannot be shared across threads")


	# Module-level state to track handler calls
	_handler_state: dict[str, Any] = {}


	def _reset_handler_state() -> None:
	"""Reset handler state before each test."""
	global _handler_state
	_handler_state = {
	"abort_called": False,
	"handler_exception": None,
	}


	def timeout_abortable_task() -> None:
	"""Task with abort handler that exits when aborted."""
	ctx = get_context()

	@ctx.on_abort
	def on_abort() -> None:
	_handler_state["abort_called"] = True

	# Poll for abort signal
	for _ in range(50):
	if _handler_state["abort_called"]:
	return
	time.sleep(0.1)


	def timeout_handler_fails_task() -> None:
	"""Task with abort handler that throws an exception."""
	ctx = get_context()

	@ctx.on_abort
	def on_abort() -> None:
	_handler_state["abort_called"] = True
	raise ValueError("Handler crashed!")

	# Sleep longer than timeout
	time.sleep(5)


	def simple_task_with_abort() -> None:
	"""Simple task with abort handler for testing."""
	ctx = get_context()

	@ctx.on_abort
	def on_abort() -> None:
	pass


	def quick_task_with_abort() -> None:
	"""Quick task that completes before timeout."""
	ctx = get_context()

	@ctx.on_abort
	def on_abort() -> None:
	pass

	time.sleep(0.2)


	def _register_test_tasks() -> None:
	"""Register test task functions if not already registered.

	Called in setUp() to ensure tasks are registered regardless of
	whether other tests have cleared the registry.
	"""
	registrations = [
	("test_timeout_abortable", timeout_abortable_task),
	("test_timeout_handler_fails", timeout_handler_fails_task),
	("test_timeout_simple", simple_task_with_abort),
	("test_timeout_quick", quick_task_with_abort),
	]
	for name, func in registrations:
	if not TaskRegistry.is_registered(name):
	TaskRegistry.register(name, func)


	class TestTimeoutHandling(SupersetTestCase):
	"""E2E tests for task timeout functionality."""

	def setUp(self) -> None:
	"""Set up test fixtures."""
	super().setUp()
	self.login(ADMIN_USERNAME)
	_register_test_tasks()
	_reset_handler_state()

	def test_timeout_with_abort_handler_results_in_timed_out_status(self) -> None:
	"""Task with timeout and abort handler should end with TIMED_OUT status."""
	_skip_if_sqlite()

	# Create task with timeout
	task_obj = TaskDAO.create_task(
	task_type="test_timeout_abortable",
	task_key=f"test_key_{uuid.uuid4().hex[:8]}",
	task_name="Test Timeout",
	scope=TaskScope.SYSTEM,
	properties={"timeout": 1}, # 1 second timeout
	)

	# Execute task via Celery executor (synchronously)
	# Use str(uuid) since Celery serializes args as JSON strings
	result = execute_task.apply(
	args=[str(task_obj.uuid), "test_timeout_abortable", (), {}]
	)

	# Verify execution completed
	assert result.successful()
	assert result.result["status"] == TaskStatus.TIMED_OUT.value

	# Verify abort handler was called
	assert _handler_state["abort_called"]

	def test_user_abort_results_in_aborted_status(self) -> None:
	"""User-initiated abort on pending task should result in ABORTED."""
	# Create task (pending state)
	task_obj = TaskDAO.create_task(
	task_type="test_timeout_simple",
	task_key=f"test_key_{uuid.uuid4().hex[:8]}",
	task_name="Test Abort Task",
	scope=TaskScope.SYSTEM,
	)

	# Cancel before execution (pending task abort)
	CancelTaskCommand(task_obj.uuid, force=True).run()

	# Refresh from DB
	db.session.expire_all()
	task_obj = db.session.query(Task).filter_by(uuid=task_obj.uuid).first()
	assert task_obj.status == TaskStatus.ABORTED.value

	def test_no_timeout_when_not_configured(self) -> None:
	"""Task without timeout should run to completion regardless of duration."""
	task_obj = TaskDAO.create_task(
	task_type="test_timeout_quick",
	task_key=f"test_key_{uuid.uuid4().hex[:8]}",
	task_name="Test No Timeout",
	scope=TaskScope.SYSTEM,
	# No timeout property
	)

	# Use str(uuid) since Celery serializes args as JSON strings
	result = execute_task.apply(
	args=[str(task_obj.uuid), "test_timeout_quick", (), {}]
	)

	assert result.successful()
	assert result.result["status"] == TaskStatus.SUCCESS.value

	def test_abort_handler_exception_results_in_failure(self) -> None:
	"""If abort handler throws during timeout, task should be FAILURE."""
	_skip_if_sqlite()

	task_obj = TaskDAO.create_task(
	task_type="test_timeout_handler_fails",
	task_key=f"test_key_{uuid.uuid4().hex[:8]}",
	task_name="Test Handler Fails",
	scope=TaskScope.SYSTEM,
	properties={"timeout": 1}, # 1 second timeout
	)

	# Use str(uuid) since Celery serializes args as JSON strings
	result = execute_task.apply(
	args=[str(task_obj.uuid), "test_timeout_handler_fails", (), {}]
	)

	assert result.successful()
	assert result.result["status"] == TaskStatus.FAILURE.value
	assert _handler_state["abort_called"]