def test_create_feature_set(db: DBInterface, db_session: Session): name = "dummy" feature_set = _create_feature_set(name) project = "proj_test" feature_set = schemas.FeatureSet(**feature_set) db.create_feature_set(db_session, project, feature_set, versioned=True) db.get_feature_set(db_session, project, name) feature_set_res = db.list_feature_sets(db_session, project) assert len(feature_set_res.feature_sets) == 1 features_res = db.list_features(db_session, project, "time") assert len(features_res.features) == 1
def _create_resources_of_all_kinds(db: DBInterface, db_session: Session, project: str): # Create several functions with several tags labels = { "name": "value", "name2": "value2", } function = { "bla": "blabla", "metadata": { "labels": labels }, "status": { "bla": "blabla" }, } function_names = ["function_name_1", "function_name_2", "function_name_3"] function_tags = ["some_tag", "some_tag2", "some_tag3"] for function_name in function_names: for function_tag in function_tags: db.store_function( db_session, function, function_name, project, tag=function_tag, versioned=True, ) # Create several artifacts with several tags artifact = { "bla": "blabla", "labels": labels, "status": { "bla": "blabla" }, } artifact_keys = ["artifact_key_1", "artifact_key_2", "artifact_key_3"] artifact_uids = ["some_uid", "some_uid2", "some_uid3"] artifact_tags = ["some_tag", "some_tag2", "some_tag3"] for artifact_key in artifact_keys: for artifact_uid in artifact_uids: for artifact_tag in artifact_tags: for artifact_iter in range(3): db.store_artifact( db_session, artifact_key, artifact, artifact_uid, artifact_iter, artifact_tag, project, ) # Create several runs run = { "bla": "blabla", "metadata": { "labels": labels }, "status": { "bla": "blabla" }, } run_uids = ["some_uid", "some_uid2", "some_uid3"] for run_uid in run_uids: for run_iter in range(3): db.store_run(db_session, run, run_uid, project, run_iter) # Create several logs log = b"some random log" log_uids = ["some_uid", "some_uid2", "some_uid3"] for log_uid in log_uids: db.store_log(db_session, log_uid, project, log) # Create several schedule schedule = { "bla": "blabla", "status": { "bla": "blabla" }, } schedule_cron_trigger = schemas.ScheduleCronTrigger(year=1999) schedule_names = ["schedule_name_1", "schedule_name_2", "schedule_name_3"] for schedule_name in schedule_names: db.create_schedule( db_session, project, schedule_name, schemas.ScheduleKinds.job, schedule, schedule_cron_trigger, labels, ) feature_set = schemas.FeatureSet( metadata=schemas.ObjectMetadata(name="dummy", tag="latest", labels={"owner": "nobody"}), spec=schemas.FeatureSetSpec( entities=[ schemas.Entity(name="ent1", value_type="str", labels={"label": "1"}) ], features=[ schemas.Feature(name="feat1", value_type="str", labels={"label": "1"}) ], ), status={}, ) db.create_feature_set(db_session, project, feature_set)