async def test_get_schedule_failure_not_found(db: Session, scheduler: Scheduler): schedule_name = "schedule-name" project = config.default_project with pytest.raises(mlrun.errors.MLRunNotFoundError) as excinfo: scheduler.get_schedule(db, project, schedule_name) assert "Schedule not found" in str(excinfo.value)
async def test_invoke_schedule( db: Session, scheduler: Scheduler, k8s_secrets_mock: tests.api.conftest.K8sSecretsMock, ): cron_trigger = schemas.ScheduleCronTrigger(year=1999) schedule_name = "schedule-name" project = config.default_project scheduled_object = _create_mlrun_function_and_matching_scheduled_object( db, project) runs = get_db().list_runs(db, project=project) assert len(runs) == 0 scheduler.create_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, schemas.ScheduleKinds.job, scheduled_object, cron_trigger, ) runs = get_db().list_runs(db, project=project) assert len(runs) == 0 response_1 = await scheduler.invoke_schedule(db, mlrun.api.schemas.AuthInfo(), project, schedule_name) runs = get_db().list_runs(db, project=project) assert len(runs) == 1 response_2 = await scheduler.invoke_schedule(db, mlrun.api.schemas.AuthInfo(), project, schedule_name) runs = get_db().list_runs(db, project=project) assert len(runs) == 2 for run in runs: assert run["status"]["state"] == RunStates.completed response_uids = [ response["data"]["metadata"]["uid"] for response in [response_1, response_2] ] db_uids = [run["metadata"]["uid"] for run in runs] assert DeepDiff( response_uids, db_uids, ignore_order=True, ) == {} schedule = scheduler.get_schedule(db, project, schedule_name, include_last_run=True) assert schedule.last_run is not None assert schedule.last_run["metadata"]["uid"] == response_uids[-1] assert schedule.last_run["metadata"]["project"] == project
def _assert_schedule_get_and_list_credentials_enrichment( db: Session, scheduler: Scheduler, project: str, schedule_name: str, expected_access_key: str, ): schedule = scheduler.get_schedule( db, project, schedule_name, include_credentials=True, ) assert schedule.credentials.access_key == expected_access_key schedules = scheduler.list_schedules(db, project, schedule_name, include_credentials=True) assert schedules.schedules[0].credentials.access_key == expected_access_key
async def test_get_schedule_datetime_fields_timezone(db: Session, scheduler: Scheduler): cron_trigger = schemas.ScheduleCronTrigger(minute="*/10") schedule_name = "schedule-name" project = config.default_project scheduler.create_schedule( db, project, schedule_name, schemas.ScheduleKinds.local_function, do_nothing, cron_trigger, ) schedule = scheduler.get_schedule(db, project, schedule_name) assert schedule.creation_time.tzinfo is not None assert schedule.next_run_time.tzinfo is not None schedules = scheduler.list_schedules(db, project) assert len(schedules.schedules) == 1 assert schedules.schedules[0].creation_time.tzinfo is not None assert schedules.schedules[0].next_run_time.tzinfo is not None
async def test_update_schedule(db: Session, scheduler: Scheduler): labels_1 = { "label1": "value1", "label2": "value2", } labels_2 = { "label3": "value3", "label4": "value4", } inactive_cron_trigger = schemas.ScheduleCronTrigger(year="1999") schedule_name = "schedule-name" project = config.default_project scheduled_object = _create_mlrun_function_and_matching_scheduled_object( db, project) runs = get_db().list_runs(db, project=project) assert len(runs) == 0 scheduler.create_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, schemas.ScheduleKinds.job, scheduled_object, inactive_cron_trigger, labels=labels_1, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, labels_1, ) # update labels scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, labels=labels_2, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, labels_2, ) # update nothing scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, labels_2, ) # update labels to empty dict scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, labels={}, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, {}, ) # update it so it runs now = datetime.now() now_plus_1_second = now + timedelta(seconds=1) now_plus_2_second = now + timedelta(seconds=2) # this way we're leaving ourselves one second to create the schedule preventing transient test failure cron_trigger = schemas.ScheduleCronTrigger( second="*/1", start_date=now_plus_1_second, end_date=now_plus_2_second, ) scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, cron_trigger=cron_trigger, ) schedule = scheduler.get_schedule(db, project, schedule_name) next_run_time = datetime( year=now_plus_2_second.year, month=now_plus_2_second.month, day=now_plus_2_second.day, hour=now_plus_2_second.hour, minute=now_plus_2_second.minute, second=now_plus_2_second.second, tzinfo=tzlocal(), ) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, cron_trigger, next_run_time, {}, ) await asyncio.sleep(2) runs = get_db().list_runs(db, project=project) assert len(runs) == 1 assert runs[0]["status"]["state"] == RunStates.completed
async def test_get_schedule(db: Session, scheduler: Scheduler): labels_1 = { "label1": "value1", "label2": "value2", } cron_trigger = schemas.ScheduleCronTrigger(year="1999") schedule_name = "schedule-name" project = config.default_project scheduler.create_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, schemas.ScheduleKinds.local_function, do_nothing, cron_trigger, labels_1, ) schedule = scheduler.get_schedule(db, project, schedule_name) # no next run time cause we put year=1999 _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.local_function, cron_trigger, None, labels_1, ) labels_2 = { "label3": "value3", "label4": "value4", } year = 2050 cron_trigger_2 = schemas.ScheduleCronTrigger(year=year, timezone="utc") schedule_name_2 = "schedule-name-2" scheduler.create_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name_2, schemas.ScheduleKinds.local_function, do_nothing, cron_trigger_2, labels_2, ) schedule_2 = scheduler.get_schedule(db, project, schedule_name_2) year_datetime = datetime(year=year, month=1, day=1, tzinfo=timezone.utc) _assert_schedule( schedule_2, project, schedule_name_2, schemas.ScheduleKinds.local_function, cron_trigger_2, year_datetime, labels_2, ) schedules = scheduler.list_schedules(db) assert len(schedules.schedules) == 2 _assert_schedule( schedules.schedules[0], project, schedule_name, schemas.ScheduleKinds.local_function, cron_trigger, None, labels_1, ) _assert_schedule( schedules.schedules[1], project, schedule_name_2, schemas.ScheduleKinds.local_function, cron_trigger_2, year_datetime, labels_2, ) schedules = scheduler.list_schedules(db, labels="label3=value3") assert len(schedules.schedules) == 1 _assert_schedule( schedules.schedules[0], project, schedule_name_2, schemas.ScheduleKinds.local_function, cron_trigger_2, year_datetime, labels_2, )
async def test_update_schedule( db: Session, scheduler: Scheduler, k8s_secrets_mock: tests.api.conftest.K8sSecretsMock, ): labels_1 = { "label1": "value1", "label2": "value2", } labels_2 = { "label3": "value3", "label4": "value4", } inactive_cron_trigger = schemas.ScheduleCronTrigger(year="1999") schedule_name = "schedule-name" project = config.default_project scheduled_object = _create_mlrun_function_and_matching_scheduled_object( db, project) runs = get_db().list_runs(db, project=project) assert len(runs) == 0 scheduler.create_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, schemas.ScheduleKinds.job, scheduled_object, inactive_cron_trigger, labels=labels_1, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, labels_1, ) # update labels scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, labels=labels_2, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, labels_2, ) # update nothing scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, labels_2, ) # update labels to empty dict scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, labels={}, ) schedule = scheduler.get_schedule(db, project, schedule_name) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, inactive_cron_trigger, None, {}, ) # update it so it runs expected_call_counter = 1 start_date, end_date = _get_start_and_end_time_for_scheduled_trigger( number_of_jobs=expected_call_counter, seconds_interval=1) # this way we're leaving ourselves one second to create the schedule preventing transient test failure cron_trigger = schemas.ScheduleCronTrigger( second="*/1", start_date=start_date, end_date=end_date, ) scheduler.update_schedule( db, mlrun.api.schemas.AuthInfo(), project, schedule_name, cron_trigger=cron_trigger, ) schedule = scheduler.get_schedule(db, project, schedule_name) next_run_time = datetime( year=end_date.year, month=end_date.month, day=end_date.day, hour=end_date.hour, minute=end_date.minute, second=end_date.second, tzinfo=tzlocal(), ) _assert_schedule( schedule, project, schedule_name, schemas.ScheduleKinds.job, cron_trigger, next_run_time, {}, ) time_to_sleep = (end_date - datetime.now()).total_seconds() + schedule_end_time_margin await asyncio.sleep(time_to_sleep) runs = get_db().list_runs(db, project=project) assert len(runs) == 1 assert runs[0]["status"]["state"] == RunStates.completed