villebro commented on code in PR #36368: URL: https://github.com/apache/superset/pull/36368#discussion_r2748366878
########## tests/unit_tests/daos/test_tasks.py: ########## @@ -0,0 +1,266 @@ +# 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. + +from unittest.mock import MagicMock, patch + +import pytest +from superset_core.api.tasks import TaskStatus + +from superset.daos.tasks import TaskDAO +from superset.models.tasks import Task + + +class TestTaskDAO: + """Test TaskDAO functionality""" + + @patch("superset.utils.core.get_user_id") + @patch("superset.daos.tasks.db.session") + def test_find_by_task_key_active(self, mock_session, mock_get_user_id): + """Test finding active task by task_key""" + mock_get_user_id.return_value = 1 + mock_task = MagicMock(spec=Task) + mock_task.task_key = "test-id" + mock_task.task_type = "test_type" + mock_task.status = TaskStatus.PENDING.value + + mock_query = MagicMock() + mock_session.query.return_value = mock_query + mock_filter = MagicMock() + mock_query.filter.return_value = mock_filter + mock_filter.one_or_none.return_value = mock_task + + result = TaskDAO.find_by_task_key("test_type", "test-key") + + assert result == mock_task + mock_session.query.assert_called_once_with(Task) + + @patch("superset.utils.core.get_user_id") + @patch("superset.daos.tasks.db.session") + def test_find_by_task_key_not_found(self, mock_session, mock_get_user_id): + """Test finding task by task_key returns None when not found""" + mock_get_user_id.return_value = 1 + mock_query = MagicMock() + mock_session.query.return_value = mock_query + mock_filter = MagicMock() + mock_query.filter.return_value = mock_filter + mock_filter.one_or_none.return_value = None + + result = TaskDAO.find_by_task_key("test_type", "nonexistent-key") + + assert result is None + + @patch("superset.utils.core.get_user_id") + @patch("superset.daos.tasks.db.session") + def test_find_by_task_key_ignores_finished_tasks( + self, mock_session, mock_get_user_id + ): + """Test that find_by_task_key only returns pending/in-progress tasks""" + mock_get_user_id.return_value = 1 + mock_query = MagicMock() + mock_session.query.return_value = mock_query + mock_filter = MagicMock() + mock_query.filter.return_value = mock_filter + mock_filter.one_or_none.return_value = None + + # Should not find SUCCESS task + result = TaskDAO.find_by_task_key("test_type", "finished-key") + assert result is None + + @patch("superset.utils.core.get_user_id") + @patch("superset.daos.tasks.TaskDAO.create") + @patch("superset.daos.tasks.db.session") + def test_create_task_success(self, mock_session, mock_create, mock_get_user_id): + """Test successful task creation. + + TaskDAO.create_task is now a pure data operation - it assumes the caller + (SubmitTaskCommand) has already checked for duplicates and holds the lock. + """ + mock_get_user_id.return_value = 1 + mock_task = MagicMock(spec=Task) + mock_task.id = 1 + mock_task.task_key = "new-task" + mock_task.task_type = "test_type" + mock_create.return_value = mock_task + + result = TaskDAO.create_task( + task_type="test_type", + task_key="new-task", + ) + + assert result == mock_task + mock_create.assert_called_once() + mock_session.commit.assert_called_once() + + @patch("superset.utils.core.get_user_id") + @patch("superset.daos.tasks.TaskDAO.create") + @patch("superset.daos.tasks.db.session") + def test_create_task_with_user_id( + self, mock_session, mock_create, mock_get_user_id + ): + """Test task creation with explicit user_id. + + TaskDAO.create_task should use the provided user_id for the task. + """ + mock_get_user_id.return_value = 99 # Different from provided user_id + mock_task = MagicMock(spec=Task) + mock_task.id = 1 + mock_task.task_key = "new-task" + mock_task.task_type = "test_type" + mock_create.return_value = mock_task + + result = TaskDAO.create_task( + task_type="test_type", + task_key="new-task", + user_id=42, # Explicit user_id + ) + + assert result == mock_task + # Verify create was called with user_id in attributes + call_args = mock_create.call_args + assert call_args[1]["attributes"]["user_id"] == 42 + + @patch("superset.utils.core.get_user_id") + @patch("superset.daos.tasks.TaskDAO.create") + @patch("superset.daos.tasks.db.session") + def test_create_task_with_properties( + self, mock_session, mock_create, mock_get_user_id + ): + """Test task creation with properties. + + Properties are set via update_properties() after task creation. + """ + mock_get_user_id.return_value = 1 + mock_task = MagicMock(spec=Task) + mock_task.id = 1 + mock_task.task_key = "new-task" + mock_task.task_type = "test_type" + mock_create.return_value = mock_task + + result = TaskDAO.create_task( + task_type="test_type", + task_key="new-task", + properties={"timeout": 300}, + ) + + assert result == mock_task + mock_task.update_properties.assert_called_once_with({"timeout": 300}) + + @patch("superset.daos.tasks.TaskDAO.find_one_or_none") + @patch("superset.daos.tasks.db.session") + def test_abort_task_pending_success(self, mock_session, mock_find): + """Test successful abort of pending task - goes directly to ABORTED""" + mock_task = MagicMock(spec=Task) + mock_task.status = TaskStatus.PENDING.value + mock_task.is_shared = False + mock_task.subscriber_count = 0 + mock_find.return_value = mock_task + + result = TaskDAO.abort_task("test-uuid") + + # Now returns Task instead of bool + assert result is mock_task + mock_task.set_status.assert_called_once_with(TaskStatus.ABORTED) + mock_session.commit.assert_called_once() + + @patch("superset.daos.tasks.TaskDAO.find_one_or_none") + @patch("superset.daos.tasks.db.session") + def test_abort_task_in_progress_abortable(self, mock_session, mock_find): + """Test abort of in-progress task with abort handler. + + Should transition to ABORTING status. + """ + mock_task = MagicMock(spec=Task) + mock_task.status = TaskStatus.IN_PROGRESS.value + mock_task.properties = {"is_abortable": True} # Dict, not MagicMock + mock_task.is_shared = False + mock_task.subscriber_count = 0 + mock_find.return_value = mock_task + + result = TaskDAO.abort_task("test-uuid") + + # Now returns Task instead of bool + assert result is mock_task + # Should set status to ABORTING, not ABORTED + assert mock_task.status == TaskStatus.ABORTING.value + mock_session.merge.assert_called_once_with(mock_task) + mock_session.commit.assert_called_once() Review Comment: This was a bug, well caught -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
