|
7 | 7 | from time import sleep |
8 | 8 | from typing import Callable, Union |
9 | 9 | from unittest import mock |
| 10 | +import datetime |
10 | 11 |
|
11 | 12 | import dirty_equals |
12 | 13 | import geopandas |
@@ -554,6 +555,7 @@ def start_job(row, connection_provider, connection, **kwargs): |
554 | 555 | 12 * 60 * 60, |
555 | 556 | "finished", |
556 | 557 | ), |
| 558 | +
|
557 | 559 | ], |
558 | 560 | ) |
559 | 561 | def test_automatic_cancel_of_too_long_running_jobs( |
@@ -645,6 +647,80 @@ def test_status_logging(self, tmp_path, job_manager, job_manager_root_dir, sleep |
645 | 647 | assert needle.search(caplog.text) |
646 | 648 |
|
647 | 649 |
|
| 650 | + @pytest.mark.parametrize( |
| 651 | + ["create_time", "start_time", "running_start_time", "end_time", "end_status", "cancel_after_seconds"], |
| 652 | + [ |
| 653 | + # Scenario 1: Missing running_start_time (None) |
| 654 | + ( |
| 655 | + "2024-09-01T09:00:00Z", # Job creation time |
| 656 | + "2024-09-01T09:00:00Z", # Job start time (should be 1 hour after create_time) |
| 657 | + None, # Missing running_start_time |
| 658 | + "2024-09-01T20:00:00Z", # Job end time |
| 659 | + "finished", # Job final status |
| 660 | + 6 * 60 * 60, # Cancel after 6 hours |
| 661 | + ), |
| 662 | + # Scenario 2: NaN running_start_time |
| 663 | + ( |
| 664 | + "2024-09-01T09:00:00Z", |
| 665 | + "2024-09-01T09:00:00Z", |
| 666 | + float("nan"), # NaN running_start_time |
| 667 | + "2024-09-01T20:00:00Z", # Job end time |
| 668 | + "finished", # Job final status |
| 669 | + 6 * 60 * 60, # Cancel after 6 hours |
| 670 | + ), |
| 671 | + ] |
| 672 | + ) |
| 673 | + def test_ensure_running_start_time_is_datetime( |
| 674 | + self, |
| 675 | + tmp_path, |
| 676 | + time_machine, |
| 677 | + create_time, |
| 678 | + start_time, |
| 679 | + running_start_time, |
| 680 | + end_time, |
| 681 | + end_status, |
| 682 | + cancel_after_seconds, |
| 683 | + dummy_backend_foo, |
| 684 | + job_manager_root_dir, |
| 685 | + ): |
| 686 | + def get_status(job_id, current_status): |
| 687 | + if rfc3339.utcnow() < start_time: |
| 688 | + return "queued" |
| 689 | + elif rfc3339.utcnow() < end_time: |
| 690 | + return "running" |
| 691 | + return end_status |
| 692 | + |
| 693 | + # Set the job status updater function for the mock backend |
| 694 | + dummy_backend_foo.job_status_updater = get_status |
| 695 | + |
| 696 | + job_manager = MultiBackendJobManager( |
| 697 | + root_dir=job_manager_root_dir, cancel_running_job_after=cancel_after_seconds |
| 698 | + ) |
| 699 | + job_manager.add_backend("foo", connection=dummy_backend_foo.connection) |
| 700 | + |
| 701 | + # Create a DataFrame representing the job database |
| 702 | + df = pd.DataFrame({ |
| 703 | + "year": [2024], |
| 704 | + "running_start_time": [running_start_time], # Initial running_start_time |
| 705 | + }) |
| 706 | + |
| 707 | + # Move the time machine to the job creation time |
| 708 | + time_machine.move_to(create_time) |
| 709 | + |
| 710 | + job_db_path = tmp_path / "jobs.csv" |
| 711 | + |
| 712 | + # Mock sleep() to skip one hour at a time instead of actually sleeping |
| 713 | + with mock.patch.object(openeo.extra.job_management.time, "sleep", new=lambda s: time_machine.shift(60 * 60)): |
| 714 | + job_manager.run_jobs(df=df, start_job=self._create_year_job, job_db=job_db_path) |
| 715 | + |
| 716 | + final_df = CsvJobDatabase(job_db_path).read() |
| 717 | + |
| 718 | + # Validate running_start_time is a valid datetime object |
| 719 | + filled_running_start_time = final_df.iloc[0]["running_start_time"] |
| 720 | + assert isinstance(rfc3339.parse_datetime(filled_running_start_time), datetime.datetime) |
| 721 | + |
| 722 | + |
| 723 | + |
648 | 724 | JOB_DB_DF_BASICS = pd.DataFrame( |
649 | 725 | { |
650 | 726 | "numbers": [3, 2, 1], |
|
0 commit comments