import os import copy import json import unittest from unittest.mock import patch from consts.consts import TaskType from models.base import with_session from utils.testcase import create_data_summary, org_id from tasks.eda import task_gen_db_schema, EDATaskJob, task_chat2data, task_suggest_questions, task_text2sql from models.edatask import EDATaskJob as OriginalEDATaskJob from models.context import Context from svc.data_explore import ThreadPoolWithLocalExecutor, setup_data_context # get "data.json" absolute path in current directory data_json_path = os.path.join(os.path.dirname(__file__), "data.json") inited_data_context = json.load(open(data_json_path, "r")) def init_data_context(data_context_id): with with_session() as s: context = Context.get_by_id(s, data_context_id) assert context is not None db_context = context.load_data() db_context.update(inited_data_context) context.update_data(db_context) class Chat2QueryTasksTests(unittest.TestCase): def setUp(self) -> None: return super().setUp() def tearDown(self) -> None: return super().tearDown() @patch("tasks.eda.setup_data_context") @patch.object(EDATaskJob, "get_by_job_id") @patch("tasks.eda.task_gen_db_description") @patch("tasks.eda.task_gen_entity") @patch("tasks.eda.task_gen_table_relations") @patch("tasks.eda.task_gen_table_description") @patch("tasks.eda.task_gen_column_description") def test_gen_db_schema( self, mocked_task_gen_column_description, mocked_task_gen_table_description, mocked_task_gen_table_relations, mocked_task_gen_entity, mocked_task_gen_db_description, mocked_eda_task_job_get_by_job_id, mocked_setup_data_context, ): mocked_setup_data_context.return_value = copy.deepcopy(inited_data_context) response = create_data_summary() assert response.status_code == 200, response.text resp_json = json.loads(response.data) data_summary_id = resp_json["result"]["data_summary_id"] job_id = resp_json["result"]["job_id"] with with_session() as s: edatask_job = OriginalEDATaskJob.get_by_job_id(s, job_id) mocked_eda_task_job_get_by_job_id.return_value = edatask_job task_gen_db_schema(data_summary_id) assert mocked_setup_data_context.called assert mocked_task_gen_column_description.called assert mocked_task_gen_table_description.called assert mocked_task_gen_table_relations.called assert mocked_task_gen_entity.called assert mocked_task_gen_db_description.called @patch.object(EDATaskJob, "get_by_job_id") @patch("svc.eda_task.ai_completion") def test_gen_suggest_questions(self, mocked_ai_completion, mocked_eda_task_job_get_by_job_id): from handlers.edatask import gen_new_session_context mocked_ai_completion.return_value = [{ "type": "type", "question": "question", "keyword": "keyword", }] response = create_data_summary() assert response.status_code == 200, response.text resp_json = json.loads(response.data) data_summary_id = resp_json["result"]["data_summary_id"] job_id = resp_json["result"]["job_id"] session_context_id = None init_data_context(data_summary_id) with with_session() as s: context = Context.get_by_id(s, data_summary_id) assert context is not None db_context = context.load_data() db_context["entity"] = {"hello": { "name": "name", "summary": "summary", "attributes": [1, 2, 3], "involved_tables": ["involved_tables"], }} context.update_data(db_context) session_context = gen_new_session_context(s, org_id, context, 1, "test") edatask_job = OriginalEDATaskJob.get_by_job_id(s, job_id) mocked_eda_task_job_get_by_job_id.return_value = edatask_job session_context_id = session_context.id assert session_context_id is not None app_id = 1 public_key = "public_key" user_token = "" task_suggest_questions(session_context_id, data_summary_id, app_id, public_key, user_token) assert mocked_ai_completion.called @patch.object(EDATaskJob, "get_by_job_id") @patch("svc.eda_task.ai_completion") def test_text2sql(self, mocked_ai_completion, mocked_eda_task_job_get_by_job_id): from handlers.edatask import gen_new_session_context mocked_ai_completion.return_value = { "sql": "SELECT * FROM table", "rewritten_sql": "SELECT * FROM table", } response = create_data_summary() assert response.status_code == 200, response.text resp_json = json.loads(response.data) data_summary_id = resp_json["result"]["data_summary_id"] job_id = resp_json["result"]["job_id"] session_context_id = None task_id, question_id = "task_id", "question_id" init_data_context(data_summary_id) with with_session() as s: context = Context.get_by_id(s, data_summary_id) assert context is not None db_context = context.load_data() db_context["entity"] = {"hello": { "name": "name", "summary": "summary", "attributes": [1, 2, 3], "involved_tables": ["involved_tables"], }} context.update_data(db_context) session_context = gen_new_session_context(s, org_id, context, 1, "test") assert session_context is not None session_context_data = session_context.load_data() session_context_data["tasks_tree"] = { question_id: { task_id: { "task": "blabla", "type": TaskType.DATA_RETRIEVAL, "sequence_no": 0, "task_id": task_id, "description": "description", "clarified_task": "clarified_task", } } } session_context.update_data(session_context_data) edatask_job = OriginalEDATaskJob.get_by_job_id(s, job_id) mocked_eda_task_job_get_by_job_id.return_value = edatask_job session_context_id = session_context.id assert session_context_id is not None task_text2sql(data_summary_id, session_context_id, task_id, question_id) assert mocked_ai_completion.called @unittest.skip("Work in progress...") @patch.object(EDATaskJob, "get_by_job_id") @patch("svc.eda_task.ai_completion") def test_question_breakdown(self, mocked_ai_completion, mocked_eda_task_job_get_by_job_id): from handlers.edatask import gen_new_session_context mocked_ai_completion.return_value = { "sql": "SELECT * FROM table", "rewritten_sql": "SELECT * FROM table", } response = create_data_summary() assert response.status_code == 200, response.text resp_json = json.loads(response.data) data_summary_id = resp_json["result"]["data_summary_id"] job_id = resp_json["result"]["job_id"] session_context_id = None task_id, question_id = "task_id", "question_id" init_data_context(data_summary_id) with with_session() as s: context = Context.get_by_id(s, data_summary_id) assert context is not None db_context = context.load_data() db_context["status"] = "done" db_context["entity"] = {"hello": { "name": "name", "summary": "summary", "attributes": [1, 2, 3], "involved_tables": ["involved_tables"], }} context.update_data(db_context) session_context = gen_new_session_context(s, org_id, context, 1, "test") assert session_context is not None session_context_data = session_context.load_data() session_context_data["tasks_tree"] = { question_id: { task_id: { "task": "blabla", "type": "Data Retrieval", "sequence_no": 0, "parent_task_id": "parent_task_id", "task_id": task_id, "description": "description", "clarified_task": "clarified_task", } } } session_context.update_data(session_context_data) edatask_job = OriginalEDATaskJob.get_by_job_id(s, job_id) mocked_eda_task_job_get_by_job_id.return_value = edatask_job session_context_id = session_context.id assert session_context_id is not None raw_question, direct_resolve, auto_run, sql_to_data, data_to_chart = "hello world?", True, True, False, False app_id, public_key, user_token = "1", "public_key", "" task_chat2data(data_summary_id, session_context_id, raw_question, direct_resolve, auto_run, sql_to_data, data_to_chart, "", app_id, public_key, user_token) assert mocked_ai_completion.called @unittest.skip("Work in progress...") @patch.object(ThreadPoolWithLocalExecutor, "submit") @patch("svc.data_explore.get_domain_column_description") @patch.object(EDATaskJob, "get_by_job_id") @patch("svc.eda_task.ai_completion") def test_gen_db_schema_full( self, mocked_ai_completion, mocked_eda_task_job_get_by_job_id, mocked_get_domain_column_description, mocked_submit, ): from handlers.edatask import gen_new_session_context mocked_ai_completion.return_value = { "sql": "SELECT * FROM table", "rewritten_sql": "SELECT * FROM table", } mocked_get_domain_column_description.return_value = "get_domain_column_description" response = create_data_summary() assert response.status_code == 200, response.text resp_json = json.loads(response.data) data_summary_id = resp_json["result"]["data_summary_id"] job_id = resp_json["result"]["job_id"] init_data_context(data_summary_id) with with_session() as s: context = Context.get_by_id(s, data_summary_id) assert context is not None db_context = context.load_data() db_context["entity"] = {"hello": { "name": "name", "summary": "summary", "attributes": [1, 2, 3], "involved_tables": ["involved_tables"], }} context.update_data(db_context) edatask_job = OriginalEDATaskJob.get_by_job_id(s, job_id) mocked_eda_task_job_get_by_job_id.return_value = edatask_job task_gen_db_schema(data_summary_id) assert mocked_ai_completion.called # ==================== Qwen Integration Tests ==================== @patch.dict('os.environ', {'QWEN_ENABLED': 'true'}) @patch("svc.eda_task.get_cluster_info") @patch("svc.eda_task.context.load_dataset_context") @patch("svc.eda_task.context.load_session_context") def test_determine_model_alibaba_cloud(self, mock_load_session, mock_load_context, mock_get_cluster_info): """Test Alibaba Cloud cluster automatically switches to Qwen model""" from svc.eda_task import EDATask from consts.consts import Claude3, Qwen # Setup mock - EDATask __init__ requires session_context mock_load_session.return_value = {"tasks_to_be_resolved": []} # Setup mock - Alibaba Cloud cluster mock_load_context.return_value = {"cluster_id": 123} mock_get_cluster_info.return_value = { "cluster": { "cloud_provider_name": "alibaba" # Alibaba Cloud } } # Create EDATask instance task = EDATask("test-session-id") # Call cloud platform detection method result_model = task._determine_model_by_cloud_provider(Claude3.SONNET) # Verify result - should be switched to Qwen.DEFAULT self.assertEqual(result_model, Qwen.DEFAULT) mock_get_cluster_info.assert_called_once_with(123) @patch.dict('os.environ', {'QWEN_ENABLED': 'true'}) @patch("svc.eda_task.get_cluster_info") @patch("svc.eda_task.context.load_dataset_context") @patch("svc.eda_task.context.load_session_context") def test_determine_model_aws_cloud(self, mock_load_session, mock_load_context, mock_get_cluster_info): """Test AWS cluster keeps the original model without hash-based routing""" from svc.eda_task import EDATask from consts.consts import Claude3 # Setup mock - EDATask __init__ requires session_context mock_load_session.return_value = {"tasks_to_be_resolved": []} # This cluster_id previously fell into hash bucket 0 and used Qwen. test_cluster_id = 2 # Setup mock - AWS cluster mock_load_context.return_value = {"cluster_id": test_cluster_id} mock_get_cluster_info.return_value = { "cluster": { "cloud_provider_name": "aws" # AWS } } # Create EDATask instance task = EDATask("test-session-id") # Call cloud platform detection method result_model = task._determine_model_by_cloud_provider(Claude3.SONNET) # Verify: non-AliCloud clusters outside the whitelist keep the original model. self.assertEqual(result_model, Claude3.SONNET) mock_get_cluster_info.assert_called_once_with(test_cluster_id) @patch.dict('os.environ', {'QWEN_ENABLED': 'true'}) @patch("svc.eda_task.get_cluster_info") @patch("svc.eda_task.context.load_dataset_context") @patch("svc.eda_task.context.load_session_context") def test_determine_model_aliyun_variant(self, mock_load_session, mock_load_context, mock_get_cluster_info): """Test Alibaba Cloud naming variant (aliyun)""" from svc.eda_task import EDATask from consts.consts import GPT_4O, Qwen # Setup mock - EDATask __init__ requires session_context mock_load_session.return_value = {"tasks_to_be_resolved": []} # Setup mock - using aliyun naming mock_load_context.return_value = {"cluster_id": 789} mock_get_cluster_info.return_value = { "cluster": { "cloud_provider_name": "aliyun" # Alibaba Cloud alternative naming } } # Create EDATask instance task = EDATask("test-session-id") # Call cloud platform detection method result_model = task._determine_model_by_cloud_provider(GPT_4O) # Verify result - should be switched to Qwen.DEFAULT self.assertEqual(result_model, Qwen.DEFAULT) @patch.dict('os.environ', {'QWEN_ENABLED': 'true'}) @patch("svc.eda_task.context.load_dataset_context") @patch("svc.eda_task.context.load_session_context") def test_determine_model_no_cluster_id(self, mock_load_session, mock_load_context): """Test keeping original model when cluster_id is missing""" from svc.eda_task import EDATask from consts.consts import Claude3 # Setup mock - EDATask __init__ requires session_context mock_load_session.return_value = {"tasks_to_be_resolved": []} # Setup mock - no cluster_id mock_load_context.return_value = {} # Create EDATask instance task = EDATask("test-session-id") # Call cloud platform detection method result_model = task._determine_model_by_cloud_provider(Claude3.SONNET) # Verify result - should keep original model self.assertEqual(result_model, Claude3.SONNET) @patch.dict('os.environ', {'QWEN_ENABLED': 'true'}) @patch("svc.eda_task.get_cluster_info") @patch("svc.eda_task.context.load_dataset_context") @patch("svc.eda_task.context.load_session_context") def test_determine_model_api_error(self, mock_load_session, mock_load_context, mock_get_cluster_info): """Test keeping original model when get_cluster_info fails""" from svc.eda_task import EDATask from consts.consts import Claude3 # Setup mock - EDATask __init__ requires session_context mock_load_session.return_value = {"tasks_to_be_resolved": []} # Setup mock - get_cluster_info raises exception mock_load_context.return_value = {"cluster_id": 999} mock_get_cluster_info.side_effect = Exception("API error") # Create EDATask instance task = EDATask("test-session-id") # Call cloud platform detection method - should not raise exception result_model = task._determine_model_by_cloud_provider(Claude3.SONNET) # Verify result - should keep original model self.assertEqual(result_model, Claude3.SONNET) @patch.dict('os.environ', {'QWEN_ENABLED': 'false'}) @patch("svc.eda_task.get_cluster_info") @patch("svc.eda_task.context.load_dataset_context") @patch("svc.eda_task.context.load_session_context") def test_determine_model_qwen_disabled(self, mock_load_session, mock_load_context, mock_get_cluster_info): """Test disabling Qwen routing when QWEN_ENABLED=false""" from svc.eda_task import EDATask from consts.consts import Claude3, Qwen # Setup mock - EDATask __init__ requires session_context mock_load_session.return_value = {"tasks_to_be_resolved": []} # Setup mock - Alibaba Cloud cluster (normally would use Qwen) mock_load_context.return_value = {"cluster_id": 123} mock_get_cluster_info.return_value = { "cluster": { "cloud_provider_name": "alibaba" # Alibaba Cloud } } # Create EDATask instance task = EDATask("test-session-id") # Call cloud platform detection method result_model = task._determine_model_by_cloud_provider(Claude3.SONNET) # Verify result - even for Alibaba Cloud, should keep original model when feature flag is disabled self.assertEqual(result_model, Claude3.SONNET)