async def test_create_schedule_failure_too_frequent_cron_trigger(
db: Session, scheduler: Scheduler
):
scheduler._min_allowed_interval = "10 minutes"
cases = [
{"second": "*"},
{"second": "1,2"},
{"second": "*/30"},
{"second": "30-35"},
{"second": "30-40/5"},
{"minute": "*"},
{"minute": "*"},
{"minute": "*/5"},
{"minute": "43-59"},
{"minute": "30-50/6"},
{"minute": "1,3,5"},
{"minute": "11,22,33,44,55,59"},
]
for case in cases:
cron_trigger = schemas.ScheduleCronTrigger(**case)
with pytest.raises(ValueError) as excinfo:
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
"project",
"schedule-name",
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
assert "Cron trigger too frequent. no more then one job" in str(excinfo.value)
async def test_create_schedule_mlrun_function(db: Session,
scheduler: Scheduler):
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)
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,
)
await asyncio.sleep(2)
runs = get_db().list_runs(db, project=project)
assert len(runs) == 1
assert runs[0]["status"]["state"] == RunStates.completed
expected_last_run_uri = f"{project}@{runs[0]['metadata']['uid']}#0"
schedule = get_db().get_schedule(db, project, schedule_name)
assert schedule.last_run_uri == expected_last_run_uri
async def test_not_skipping_delayed_schedules(db: Session,
scheduler: Scheduler):
global call_counter
call_counter = 0
now = datetime.now()
expected_call_counter = 1
now_plus_1_seconds = now + timedelta(seconds=1)
now_plus_2_seconds = now + timedelta(seconds=1 + expected_call_counter)
# 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_seconds,
end_date=now_plus_2_seconds)
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schemas.ScheduleKinds.local_function,
bump_counter,
cron_trigger,
)
# purposely doing time.sleep to block the reactor to ensure a job is still scheduled although its planned
# execution time passed
time.sleep(2 + expected_call_counter)
await asyncio.sleep(1)
assert call_counter == expected_call_counter
async def test_rescheduling(db: Session, scheduler: Scheduler):
global call_counter
call_counter = 0
expected_call_counter = 2
start_date, end_date = _get_start_and_end_time_for_scheduled_trigger(
number_of_jobs=expected_call_counter, seconds_interval=1)
cron_trigger = schemas.ScheduleCronTrigger(second="*/1",
start_date=start_date,
end_date=end_date)
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schemas.ScheduleKinds.local_function,
bump_counter,
cron_trigger,
)
# wait so one run will complete
time_to_sleep = (start_date - datetime.now()).total_seconds() + 1
await asyncio.sleep(time_to_sleep)
# stop the scheduler and assert indeed only one call happened
await scheduler.stop()
assert call_counter == 1
# start the scheduler and and assert another run
await scheduler.start(db)
await asyncio.sleep(1 + schedule_end_time_margin)
assert call_counter == 2
async def test_create_schedule_failure_already_exists(db: Session,
scheduler: Scheduler):
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,
)
with pytest.raises(
mlrun.errors.MLRunConflictError,
match=
rf"Conflict - Schedule already exists: {project}/{schedule_name}",
):
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
async def test_rescheduling(db: Session, scheduler: Scheduler):
global call_counter
call_counter = 0
now = datetime.now()
now_plus_2_seconds = now + timedelta(seconds=2)
cron_trigger = schemas.ScheduleCronTrigger(second="*/1",
start_date=now,
end_date=now_plus_2_seconds)
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schemas.ScheduleKinds.local_function,
bump_counter,
cron_trigger,
)
# wait so one run will complete
await asyncio.sleep(1)
# stop the scheduler and assert indeed only one call happened
await scheduler.stop()
assert call_counter == 1
# start the scheduler and and assert another run
await scheduler.start(db)
await asyncio.sleep(1)
assert call_counter == 2
async def test_create_schedule(db: Session, scheduler: Scheduler):
global call_counter
call_counter = 0
expected_call_counter = 5
start_date, end_date = _get_start_and_end_time_for_scheduled_trigger(
number_of_jobs=5, 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)
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schemas.ScheduleKinds.local_function,
bump_counter,
cron_trigger,
)
# The trigger is defined with `second="*/1"` meaning it runs on round seconds,
# but executing the actual functional code - bumping the counter - happens a few microseconds afterwards.
# To avoid transient errors on slow systems, we add extra margin.
time_to_sleep = (end_date -
datetime.now()).total_seconds() + schedule_end_time_margin
await asyncio.sleep(time_to_sleep)
assert call_counter == expected_call_counter
def _create_do_nothing_schedule(db: Session, scheduler: Scheduler,
project: str, name: str):
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
name,
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
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
async def test_schedule_crud_secrets_handling(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
mlrun.api.utils.auth.verifier.AuthVerifier(
).is_jobs_auth_required = unittest.mock.Mock(return_value=True)
for schedule_name in ["valid-secret-key", "invalid/secret/key"]:
project = config.default_project
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
access_key = "some-user-access-key"
username = "******"
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(username=username,
access_key=access_key),
project,
schedule_name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
_assert_schedule_secrets(scheduler, project, schedule_name, username,
access_key)
_assert_schedule_get_and_list_credentials_enrichment(
db, scheduler, project, schedule_name, access_key)
username = "******"
access_key = "new-access-key"
# update labels
scheduler.update_schedule(
db,
mlrun.api.schemas.AuthInfo(username=username,
access_key=access_key),
project,
schedule_name,
labels={"label-key": "label-value"},
)
_assert_schedule_secrets(scheduler, project, schedule_name, username,
access_key)
_assert_schedule_get_and_list_credentials_enrichment(
db, scheduler, project, schedule_name, access_key)
# delete schedule
scheduler.delete_schedule(
db,
project,
schedule_name,
)
_assert_schedule_secrets(scheduler, project, schedule_name, None, None)
async def test_schedule_upgrade_from_scheduler_without_credentials_store(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
"""
Continue here
this test doesn't work cause reload schedules takes for granted that there is a session, which made me think whether
session is enough - after runtimes refactor and getting auth info out of sqlrundb is will be enough, so also can
remove the mlrun.api.utils.auth.AuthVerifier().generate_auth_info_from_session call from scheduler run wrapper
"""
name = "schedule-name"
project = config.default_project
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
now = datetime.now()
expected_call_counter = 3
now_plus_2_seconds = now + timedelta(seconds=2)
now_plus_5_seconds = now + timedelta(seconds=2 + expected_call_counter)
cron_trigger = schemas.ScheduleCronTrigger(second="*/1",
start_date=now_plus_2_seconds,
end_date=now_plus_5_seconds)
# we're before upgrade so create a schedule with empty auth info
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
# stop scheduler, reconfigure to store credentials and start again (upgrade)
await scheduler.stop()
scheduler._store_schedule_credentials_in_secrets = True
await scheduler.start(db)
# at this point the schedule is inside the scheduler without auth_info, so the first trigger should try to generate
# auth info, mock the functions for this
username = "******"
session = "some-session"
mlrun.api.utils.singletons.project_member.get_project_member(
).get_project_owner = unittest.mock.Mock(
return_value=mlrun.api.schemas.ProjectOwner(username=username,
session=session))
await asyncio.sleep(2 + expected_call_counter + 1)
runs = get_db().list_runs(db, project=project)
assert len(runs) == 3
assert (mlrun.api.utils.singletons.project_member.get_project_member().
get_project_owner.call_count == 1)
async def test_schedule_upgrade_from_scheduler_without_credentials_store(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
name = "schedule-name"
project = config.default_project
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
expected_call_counter = 3
start_date, end_date = _get_start_and_end_time_for_scheduled_trigger(
number_of_jobs=expected_call_counter, seconds_interval=1)
cron_trigger = schemas.ScheduleCronTrigger(second="*/1",
start_date=start_date,
end_date=end_date)
# we're before upgrade so create a schedule with empty auth info
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
# stop scheduler, reconfigure to store credentials and start again (upgrade)
await scheduler.stop()
mlrun.api.utils.auth.verifier.AuthVerifier(
).is_jobs_auth_required = unittest.mock.Mock(return_value=True)
await scheduler.start(db)
# at this point the schedule is inside the scheduler without auth_info, so the first trigger should try to generate
# auth info, mock the functions for this
username = "******"
session = "some-session"
mlrun.api.utils.singletons.project_member.get_project_member(
).get_project_owner = unittest.mock.Mock(
return_value=mlrun.api.schemas.ProjectOwner(username=username,
session=session))
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) == 3
assert (mlrun.api.utils.singletons.project_member.get_project_member().
get_project_owner.call_count == 1)
async def test_schedule_job_concurrency_limit(
db: Session,
scheduler: Scheduler,
concurrency_limit: int,
run_amount: int,
schedule_kind: schemas.ScheduleKinds,
):
global call_counter
call_counter = 0
now = datetime.now()
now_plus_1_seconds = now + timedelta(seconds=1)
now_plus_5_seconds = now + timedelta(seconds=5)
cron_trigger = schemas.ScheduleCronTrigger(
second="*/1", start_date=now_plus_1_seconds, end_date=now_plus_5_seconds
)
schedule_name = "schedule-name"
project = config.default_project
scheduled_object = (
_create_mlrun_function_and_matching_scheduled_object(
db, project, handler="sleep_two_seconds"
)
if schedule_kind == schemas.ScheduleKinds.job
else bump_counter_and_wait
)
runs = get_db().list_runs(db, project=project)
assert len(runs) == 0
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schedule_kind,
scheduled_object,
cron_trigger,
concurrency_limit=concurrency_limit,
)
# wait so all runs will complete
await asyncio.sleep(7)
if schedule_kind == schemas.ScheduleKinds.job:
runs = get_db().list_runs(db, project=project)
assert len(runs) == run_amount
else:
assert call_counter == run_amount
async def test_rescheduling_secrets_storing(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
mlrun.api.utils.auth.verifier.AuthVerifier(
).is_jobs_auth_required = unittest.mock.Mock(return_value=True)
name = "schedule-name"
project = config.default_project
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
username = "******"
access_key = "some-user-access-key"
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(username=username, access_key=access_key),
project,
name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
jobs = scheduler._list_schedules_from_scheduler(project)
assert jobs[0].args[5].access_key == access_key
assert jobs[0].args[5].username == username
k8s_secrets_mock.assert_project_secrets(
project,
{
mlrun.api.crud.Secrets().generate_schedule_access_key_secret_key(name):
access_key,
mlrun.api.crud.Secrets().generate_schedule_username_secret_key(name):
username,
},
)
await scheduler.stop()
jobs = scheduler._list_schedules_from_scheduler(project)
assert jobs == []
await scheduler.start(db)
jobs = scheduler._list_schedules_from_scheduler(project)
assert jobs[0].args[5].username == username
assert jobs[0].args[5].access_key == access_key
async def test_schedule_access_key_generation(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
mlrun.api.utils.auth.verifier.AuthVerifier(
).is_jobs_auth_required = unittest.mock.Mock(return_value=True)
project = config.default_project
schedule_name = "schedule-name"
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
access_key = "generated-access-key"
mlrun.api.utils.auth.verifier.AuthVerifier(
).get_or_create_access_key = unittest.mock.Mock(return_value=access_key)
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
schedule_name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
mlrun.api.utils.auth.verifier.AuthVerifier(
).get_or_create_access_key.assert_called_once()
_assert_schedule_secrets(scheduler, project, schedule_name, None,
access_key)
access_key = "generated-access-key-2"
mlrun.api.utils.auth.verifier.AuthVerifier(
).get_or_create_access_key = unittest.mock.Mock(return_value=access_key)
scheduler.update_schedule(
db,
mlrun.api.schemas.AuthInfo(
access_key=mlrun.model.Credentials.generate_access_key),
project,
schedule_name,
labels={"label-key": "label-value"},
)
mlrun.api.utils.auth.verifier.AuthVerifier(
).get_or_create_access_key.assert_called_once()
_assert_schedule_secrets(scheduler, project, schedule_name, None,
access_key)
async def test_create_schedule_mlrun_function(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
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)
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,
)
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) == expected_call_counter
assert runs[0]["status"]["state"] == RunStates.completed
# the default of list_runs returns the the list descending by date.
expected_last_run_uri = f"{project}@{runs[0]['metadata']['uid']}#0"
schedule = get_db().get_schedule(db, project, schedule_name)
assert schedule.last_run_uri == expected_last_run_uri
async def test_delete_schedule(db: Session, scheduler: Scheduler):
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
project,
schedule_name,
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
schedules = scheduler.list_schedules(db)
assert len(schedules.schedules) == 1
scheduler.delete_schedule(db, project, schedule_name)
schedules = scheduler.list_schedules(db)
assert len(schedules.schedules) == 0
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_create_schedule_success_cron_trigger_validation(
db: Session, scheduler: Scheduler):
scheduler._min_allowed_interval = "10 minutes"
cases = [
{
"second": "1",
"minute": "19"
},
{
"second": "30",
"minute": "9,19"
},
{
"minute": "*/10"
},
{
"minute": "20-40/10"
},
{
"hour": "1"
},
{
"year": "1999"
},
{
"year": "2050"
},
]
for index, case in enumerate(cases):
cron_trigger = schemas.ScheduleCronTrigger(**case)
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
"project",
f"schedule-name-{index}",
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
async def test_create_schedule(db: Session, scheduler: Scheduler):
global call_counter
call_counter = 0
now = datetime.now()
expected_call_counter = 5
now_plus_1_seconds = now + timedelta(seconds=1)
now_plus_5_seconds = now + timedelta(seconds=1 + expected_call_counter)
# 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_seconds,
end_date=now_plus_5_seconds)
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
project,
schedule_name,
schemas.ScheduleKinds.local_function,
bump_counter,
cron_trigger,
)
await asyncio.sleep(1 + expected_call_counter)
assert call_counter == expected_call_counter
async def test_rescheduling_secrets_storing(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
scheduler._store_schedule_credentials_in_secrets = True
name = "schedule-name"
project = config.default_project
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
session = "some-user-session"
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(session=session),
project,
name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
jobs = scheduler._list_schedules_from_scheduler(project)
assert jobs[0].args[5].session == session
k8s_secrets_mock.assert_project_secrets(
project,
{mlrun.api.crud.Secrets().generate_schedule_secret_key(name): session})
await scheduler.stop()
jobs = scheduler._list_schedules_from_scheduler(project)
assert jobs == []
await scheduler.start(db)
jobs = scheduler._list_schedules_from_scheduler(project)
assert jobs[0].args[5].session == session
async def test_list_schedules_name_filter(db: Session, scheduler: Scheduler):
cases = [
{"name": "some_prefix-mlrun", "should_find": True},
{"name": "some_prefix-mlrun-some_suffix", "should_find": True},
{"name": "mlrun-some_suffix", "should_find": True},
{"name": "mlrun", "should_find": True},
{"name": "MLRun", "should_find": True},
{"name": "bla-MLRun-bla", "should_find": True},
{"name": "mlun", "should_find": False},
{"name": "mlurn", "should_find": False},
{"name": "mluRn", "should_find": False},
]
cron_trigger = schemas.ScheduleCronTrigger(minute="*/10")
project = config.default_project
expected_schedule_names = []
for case in cases:
name = case["name"]
should_find = case["should_find"]
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(),
project,
name,
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
if should_find:
expected_schedule_names.append(name)
schedules = scheduler.list_schedules(db, project, "~mlrun")
assert len(schedules.schedules) == len(expected_schedule_names)
for schedule in schedules.schedules:
assert schedule.name in expected_schedule_names
expected_schedule_names.remove(schedule.name)
async def test_create_schedule_failure_already_exists(db: Session,
scheduler: Scheduler):
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
schedule_name = "schedule-name"
project = config.default_project
scheduler.create_schedule(
db,
project,
schedule_name,
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
with pytest.raises(mlrun.errors.MLRunConflictError) as excinfo:
scheduler.create_schedule(
db,
project,
schedule_name,
schemas.ScheduleKinds.local_function,
do_nothing,
cron_trigger,
)
assert "Conflict - Schedule already exists" in str(excinfo.value)
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_schedule_crud_secrets_handling(
db: Session,
scheduler: Scheduler,
k8s_secrets_mock: tests.api.conftest.K8sSecretsMock,
):
scheduler._store_schedule_credentials_in_secrets = True
for schedule_name in ["valid-secret-key", "invalid/secret/key"]:
project = config.default_project
scheduled_object = _create_mlrun_function_and_matching_scheduled_object(
db, project)
session = "some-user-session"
cron_trigger = schemas.ScheduleCronTrigger(year="1999")
scheduler.create_schedule(
db,
mlrun.api.schemas.AuthInfo(session=session),
project,
schedule_name,
schemas.ScheduleKinds.job,
scheduled_object,
cron_trigger,
)
secret_key = mlrun.api.crud.Secrets().generate_schedule_secret_key(
schedule_name)
key_map_secret_key = (
mlrun.api.crud.Secrets().generate_schedule_key_map_secret_key())
secret_value = mlrun.api.crud.Secrets().get_secret(
project,
scheduler._secrets_provider,
secret_key,
allow_secrets_from_k8s=True,
allow_internal_secrets=True,
key_map_secret_key=key_map_secret_key,
)
assert secret_value == session
session = "new-session"
# update labels
scheduler.update_schedule(
db,
mlrun.api.schemas.AuthInfo(session=session),
project,
schedule_name,
labels={"label-key": "label-value"},
)
secret_value = mlrun.api.crud.Secrets().get_secret(
project,
scheduler._secrets_provider,
secret_key,
allow_secrets_from_k8s=True,
allow_internal_secrets=True,
key_map_secret_key=key_map_secret_key,
)
assert secret_value == session
# delete schedule
scheduler.delete_schedule(
db,
project,
schedule_name,
)
secret_value = mlrun.api.crud.Secrets().get_secret(
project,
scheduler._secrets_provider,
secret_key,
allow_secrets_from_k8s=True,
allow_internal_secrets=True,
key_map_secret_key=key_map_secret_key,
)
assert secret_value is None
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
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,
)
