Bill/AI-Trader
1
0
Fork 0
mirror of https://github.com/Xe138/AI-Trader.git synced 2026-04-01 17:17:24 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix(tests): update unit tests to mock _prepare_data

Browse Source 
Update existing simulation_worker unit tests to account for new _prepare_data integration:
- Mock _prepare_data to return available dates
- Update mock executors to return proper result dicts with model/date fields

Note: Some tests need additional work to properly verify job status updates.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Bill Ballou
2025-11-01 23:55:53 -04:00
parent 5e5354e2af
commit 91ffb7c71e
1 changed files with 64 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
tests/unit/test_simulation_worker.py
View File

@@ -49,10 +49,17 @@ class TestSimulationWorkerExecution:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
# Mock _prepare_data to return both dates
worker._prepare_data = Mock(return_value=(["2025-01-16", "2025-01-17"], []))
# Mock ModelDayExecutor
with patch("api.simulation_worker.ModelDayExecutor") as mock_executor_class:
mock_executor = Mock()
mock_executor.execute.return_value = {"success": True}
mock_executor.execute.return_value = {
"success": True,
"model": "test-model",
"date": "2025-01-16"
}
mock_executor_class.return_value = mock_executor
worker.run()
@@ -74,12 +81,19 @@ class TestSimulationWorkerExecution:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
# Mock _prepare_data to return both dates
worker._prepare_data = Mock(return_value=(["2025-01-16", "2025-01-17"], []))
execution_order = []
def track_execution(job_id, date, model_sig, config_path, db_path):
executor = Mock()
execution_order.append((date, model_sig))
executor.execute.return_value = {"success": True}
executor.execute.return_value = {
"success": True,
"model": model_sig,
"date": date
}
return executor
with patch("api.simulation_worker.ModelDayExecutor", side_effect=track_execution):
@@ -112,9 +126,16 @@ class TestSimulationWorkerExecution:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
# Mock _prepare_data to return the date
worker._prepare_data = Mock(return_value=(["2025-01-16"], []))
with patch("api.simulation_worker.ModelDayExecutor") as mock_executor_class:
mock_executor = Mock()
mock_executor.execute.return_value = {"success": True}
mock_executor.execute.return_value = {
"success": True,
"model": "gpt-5",
"date": "2025-01-16"
}
mock_executor_class.return_value = mock_executor
worker.run()
@@ -137,13 +158,21 @@ class TestSimulationWorkerExecution:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
# Mock _prepare_data to return the date
worker._prepare_data = Mock(return_value=(["2025-01-16"], []))
call_count = 0
def mixed_results(*args, **kwargs):
nonlocal call_count
executor = Mock()
# First model succeeds, second fails
executor.execute.return_value = {"success": call_count == 0}
success = (call_count == 0)
executor.execute.return_value = {
"success": success,
"model": kwargs.get("model_sig", "unknown"),
"date": kwargs.get("date", "2025-01-16")
}
call_count += 1
return executor
@@ -173,6 +202,9 @@ class TestSimulationWorkerErrorHandling:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
# Mock _prepare_data to return the date
worker._prepare_data = Mock(return_value=(["2025-01-16"], []))
execution_count = 0
def counting_executor(*args, **kwargs):
@@ -181,9 +213,18 @@ class TestSimulationWorkerErrorHandling:
executor = Mock()
# Second model fails
if execution_count == 2:
executor.execute.return_value = {"success": False, "error": "Model failed"}
executor.execute.return_value = {
"success": False,
"error": "Model failed",
"model": kwargs.get("model_sig", "unknown"),
"date": kwargs.get("date", "2025-01-16")
}
else:
executor.execute.return_value = {"success": True}
executor.execute.return_value = {
"success": True,
"model": kwargs.get("model_sig", "unknown"),
"date": kwargs.get("date", "2025-01-16")
}
return executor
with patch("api.simulation_worker.ModelDayExecutor", side_effect=counting_executor):
@@ -206,7 +247,9 @@ class TestSimulationWorkerErrorHandling:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
with patch("api.simulation_worker.ModelDayExecutor", side_effect=Exception("Unexpected error")):
# Mock _prepare_data to raise exception
worker._prepare_data = Mock(side_effect=Exception("Unexpected error"))
worker.run()
# Check job status
@@ -233,16 +276,27 @@ class TestSimulationWorkerConcurrency:
worker = SimulationWorker(job_id=job_id, db_path=clean_db)
# Mock _prepare_data to return the date
worker._prepare_data = Mock(return_value=(["2025-01-16"], []))
with patch("api.simulation_worker.ModelDayExecutor") as mock_executor_class:
mock_executor = Mock()
mock_executor.execute.return_value = {"success": True}
mock_executor.execute.return_value = {
"success": True,
"model": "test-model",
"date": "2025-01-16"
}
mock_executor_class.return_value = mock_executor
# Mock ThreadPoolExecutor to verify it's being used
with patch("api.simulation_worker.ThreadPoolExecutor") as mock_pool:
mock_pool_instance = Mock()
mock_pool.return_value.__enter__.return_value = mock_pool_instance
mock_pool_instance.submit.return_value = Mock(result=lambda: {"success": True})
mock_pool_instance.submit.return_value = Mock(result=lambda: {
"success": True,
"model": "test-model",
"date": "2025-01-16"
})
worker.run()