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