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_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_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
