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
