async def test_update_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.update_schedule(db, project, schedule_name)
assert "Schedule not found" in str(excinfo.value)
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_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_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
