def test_list_runs_name_filter(db: DBInterface, db_session: Session): run_name_1 = "run_name_1" run_name_2 = "run_name_2" run_1 = {"metadata": {"name": run_name_1}, "status": {"bla": "blabla"}} run_2 = {"metadata": {"name": run_name_2}, "status": {"bla": "blabla"}} # run with no name run_3 = {"metadata": {}, "status": {"bla": "blabla"}} run_uid_1 = "run_uid_1" run_uid_2 = "run_uid_2" run_uid_3 = "run_uid_3" db.store_run( db_session, run_1, run_uid_1, ) db.store_run( db_session, run_2, run_uid_2, ) db.store_run( db_session, run_3, run_uid_3, ) runs = db.list_runs(db_session) assert len(runs) == 3 runs = db.list_runs(db_session, name=run_name_1) assert len(runs) == 1 assert runs[0]["metadata"]["name"] == run_name_1 runs = db.list_runs(db_session, name=run_name_2) assert len(runs) == 1 assert runs[0]["metadata"]["name"] == run_name_2 runs = db.list_runs(db_session, name="run_name") assert len(runs) == 2
def _update_ui_url( self, db: DBInterface, db_session: Session, project: str, uid: str, crd_object, run: Dict = None, leader_session: Optional[str] = None, ): app_state = (crd_object.get("status", {}).get("applicationState", {}).get("state")) state = SparkApplicationStates.spark_application_state_to_run_state( app_state) ui_url = None if state == RunStates.running: ui_url = (crd_object.get("status", {}).get("driverInfo", {}).get("webUIIngressAddress")) db_ui_url = run.get("status", {}).get("ui_url") if db_ui_url == ui_url: return run.setdefault("status", {})["ui_url"] = ui_url db.store_run(db_session, run, uid, project, leader_session=leader_session)
def test_list_functions_no_tags(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}} function_2 = {'bla2': 'blabla', 'status': {'bla': 'blabla'}} function_name_1 = 'function_name_1' # It is impossible to create a function without tag - only to create with a tag, and then tag another function with # the same tag tag = 'some_tag' function_1_hash_key = db.store_function(db_session, function_1, function_name_1, tag=tag, versioned=True) function_2_hash_key = db.store_function(db_session, function_2, function_name_1, tag=tag, versioned=True) assert function_1_hash_key != function_2_hash_key functions = db.list_functions(db_session, function_name_1) assert len(functions) == 2 # Verify function 1 without tag and has not status for function in functions: if function['metadata']['hash'] == function_1_hash_key: assert function['metadata']['tag'] == '' assert function['status'] is None
def test_list_functions_no_tags(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla", "status": {"bla": "blabla"}} function_2 = {"bla2": "blabla", "status": {"bla": "blabla"}} function_name_1 = "function_name_1" # It is impossible to create a function without tag - only to create with a tag, and then tag another function with # the same tag tag = "some_tag" function_1_hash_key = db.store_function(db_session, function_1, function_name_1, tag=tag, versioned=True) function_2_hash_key = db.store_function(db_session, function_2, function_name_1, tag=tag, versioned=True) assert function_1_hash_key != function_2_hash_key functions = db.list_functions(db_session, function_name_1) assert len(functions) == 2 # Verify function 1 without tag and has not status for function in functions: if function["metadata"]["hash"] == function_1_hash_key: assert function["metadata"]["tag"] == "" assert function["status"] is None
def test_store_artifact_tagging(db: DBInterface, db_session: Session): artifact_1_key = "artifact_key_1" artifact_1_body = _generate_artifact(artifact_1_key) artifact_1_kind = ChartArtifact.kind artifact_1_with_kind_body = _generate_artifact(artifact_1_key, kind=artifact_1_kind) artifact_1_uid = "artifact_uid" artifact_1_with_kind_uid = "artifact_uid_2" db.store_artifact( db_session, artifact_1_key, artifact_1_body, artifact_1_uid, ) db.store_artifact( db_session, artifact_1_key, artifact_1_with_kind_body, artifact_1_with_kind_uid, ) artifact = db.read_artifact(db_session, artifact_1_key, tag="latest") assert artifact["kind"] == artifact_1_kind artifact = db.read_artifact(db_session, artifact_1_key, tag=artifact_1_uid) assert artifact.get("kind") is None artifacts = db.list_artifacts(db_session, artifact_1_key, tag="latest") assert len(artifacts) == 1 artifacts = db.list_artifacts(db_session, artifact_1_key, tag=artifact_1_uid) assert len(artifacts) == 1
def test_list_artifact_name_filter(db: DBInterface, db_session: Session): artifact_name_1 = "artifact_name_1" artifact_name_2 = "artifact_name_2" artifact_1 = _generate_artifact(artifact_name_1) artifact_2 = _generate_artifact(artifact_name_2) uid = "artifact_uid" db.store_artifact( db_session, artifact_name_1, artifact_1, uid, ) db.store_artifact( db_session, artifact_name_2, artifact_2, uid, ) artifacts = db.list_artifacts(db_session) assert len(artifacts) == 2 artifacts = db.list_artifacts(db_session, name=artifact_name_1) assert len(artifacts) == 1 assert artifacts[0]["metadata"]["name"] == artifact_name_1 artifacts = db.list_artifacts(db_session, name=artifact_name_2) assert len(artifacts) == 1 assert artifacts[0]["metadata"]["name"] == artifact_name_2 artifacts = db.list_artifacts(db_session, name="artifact_name") assert len(artifacts) == 2
def _is_runtime_resource_run_in_transient_state( self, db: DBInterface, db_session: Session, runtime_resource: Dict) -> bool: """ A runtime can have different underlying resources (like pods or CRDs) - to generalize we call it runtime resource. This function will verify whether the Run object related to this runtime resource is in transient state. This is useful in order to determine whether an object can be removed. for example, a kubejob's pod might be in completed state, but we would like to verify that the run is completed as well to verify the logs were collected before we're removing the pod. """ project, uid = self._resolve_runtime_resource_run(runtime_resource) # if no uid, assume in stable state if not uid: return False run = db.read_run(db_session, uid, project) if run.get('status', {}).get('state') not in FunctionStates.stable_phases(): return True # give some grace period now = datetime.now(timezone.utc) last_update_str = run.get('status', {}).get('last_update', now) last_update = datetime.fromisoformat(last_update_str) if last_update + timedelta(seconds=float( config.runtime_resources_deletion_grace_period)) > now: return True return False
def test_list_functions_by_tag(db: DBInterface, db_session: Session): tag = "function_name_1" names = ["some_name", "some_name2", "some_name3"] for name in names: function_body = {"metadata": {"name": name}} db.store_function(db_session, function_body, name, tag=tag, versioned=True) functions = db.list_functions(db_session, tag=tag) assert len(functions) == len(names) for function in functions: function_name = function["metadata"]["name"] names.remove(function_name) assert len(names) == 0
def test_list_functions_multiple_tags(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla", "status": {"bla": "blabla"}} function_name_1 = "function_name_1" tags = ["some_tag", "some_tag2", "some_tag3"] for tag in tags: db.store_function(db_session, function_1, function_name_1, tag=tag, versioned=True) functions = db.list_functions(db_session, function_name_1) assert len(functions) == len(tags) for function in functions: function_tag = function["metadata"]["tag"] tags.remove(function_tag) assert len(tags) == 0
def test_list_functions_multiple_tags(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}} function_name_1 = 'function_name_1' tags = ['some_tag', 'some_tag2', 'some_tag3'] for tag in tags: db.store_function(db_session, function_1, function_name_1, tag=tag, versioned=True) functions = db.list_functions(db_session, function_name_1) assert len(functions) == len(tags) for function in functions: function_tag = function['metadata']['tag'] tags.remove(function_tag) assert len(tags) == 0
def test_store_artifact_restoring_multiple_tags(db: DBInterface, db_session: Session): artifact_key = "artifact_key_1" artifact_1_uid = "artifact_uid_1" artifact_2_uid = "artifact_uid_2" artifact_1_body = _generate_artifact(artifact_key, uid=artifact_1_uid) artifact_2_body = _generate_artifact(artifact_key, uid=artifact_2_uid) artifact_1_tag = "artifact_tag_1" artifact_2_tag = "artifact_tag_2" db.store_artifact( db_session, artifact_key, artifact_1_body, artifact_1_uid, tag=artifact_1_tag, ) db.store_artifact( db_session, artifact_key, artifact_2_body, artifact_2_uid, tag=artifact_2_tag, ) artifacts = db.list_artifacts(db_session, artifact_key, tag="*") assert len(artifacts) == 2 expected_uids = [artifact_1_uid, artifact_2_uid] uids = [artifact["metadata"]["uid"] for artifact in artifacts] assert deepdiff.DeepDiff(expected_uids, uids, ignore_order=True,) == {} expected_tags = [artifact_1_tag, artifact_2_tag] tags = [artifact["tag"] for artifact in artifacts] assert deepdiff.DeepDiff(expected_tags, tags, ignore_order=True,) == {} artifact = db.read_artifact(db_session, artifact_key, tag=artifact_1_tag) assert artifact["metadata"]["uid"] == artifact_1_uid assert artifact["tag"] == artifact_1_tag artifact = db.read_artifact(db_session, artifact_key, tag=artifact_2_tag) assert artifact["metadata"]["uid"] == artifact_2_uid assert artifact["tag"] == artifact_2_tag
def test_get_function_by_hash_key(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}} function_name_1 = 'function_name_1' function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=True) function_queried_without_hash_key = db.get_function( db_session, function_name_1) assert function_queried_without_hash_key is not None # Verifying function is queryable by hash_key function_queried_with_hash_key = db.get_function( db_session, function_name_1, hash_key=function_hash_key) assert function_queried_with_hash_key is not None # function queried by hash shouldn't have tag assert function_queried_without_hash_key['metadata']['tag'] == 'latest' assert function_queried_with_hash_key['metadata']['tag'] == ''
def test_get_function_by_hash_key(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla", "status": {"bla": "blabla"}} function_name_1 = "function_name_1" function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=True) function_queried_without_hash_key = db.get_function( db_session, function_name_1) assert function_queried_without_hash_key is not None # Verifying function is queryable by hash_key function_queried_with_hash_key = db.get_function( db_session, function_name_1, hash_key=function_hash_key) assert function_queried_with_hash_key is not None # function queried by hash shouldn't have tag assert function_queried_without_hash_key["metadata"]["tag"] == "latest" assert function_queried_with_hash_key["metadata"]["tag"] == ""
def test_store_function_default_to_latest(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla'} function_name_1 = 'function_name_1' function_hash_key = db.store_function(db_session, function_1, function_name_1) assert function_hash_key is not None function_queried_without_tag = db.get_function(db_session, function_name_1) function_queried_without_tag_hash = function_queried_without_tag[ 'metadata']['hash'] assert function_hash_key == function_queried_without_tag_hash assert function_queried_without_tag['metadata']['tag'] == 'latest' function_queried_with_tag = db.get_function(db_session, function_name_1, tag='latest') function_queried_without_tag_hash = function_queried_with_tag['metadata'][ 'hash'] assert function_queried_with_tag is not None assert function_queried_with_tag['metadata']['tag'] == 'latest' assert function_queried_without_tag_hash == function_queried_without_tag_hash
def test_delete_project(db: DBInterface, db_session: Session): project_to_keep = "project_to_keep" project_to_remove = "project_to_remove" _create_resources_of_all_kinds(db, db_session, project_to_keep) _create_resources_of_all_kinds(db, db_session, project_to_remove) project_to_keep_table_name_records_count_map_before_project_removal = _assert_resources_in_project( db, db_session, project_to_keep) _assert_resources_in_project(db, db_session, project_to_remove) db.delete_project(db_session, project_to_remove) project_to_keep_table_name_records_count_map_after_project_removal = _assert_resources_in_project( db, db_session, project_to_keep) _assert_resources_in_project(db, db_session, project_to_remove, assert_no_resources=True) assert (deepdiff.DeepDiff( project_to_keep_table_name_records_count_map_before_project_removal, project_to_keep_table_name_records_count_map_after_project_removal, ignore_order=True, ) == {})
def test_get_function_by_tag(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla", "status": {"bla": "blabla"}} function_name_1 = "function_name_1" function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=True) function_queried_by_tag = db.get_function(db_session, function_name_1, tag="latest") function_queried_by_hash_key = db.get_function(db_session, function_name_1, hash_key=function_hash_key) function_not_queried_by_tag_hash = function_queried_by_hash_key[ "metadata"]["hash"] assert function_hash_key == function_not_queried_by_tag_hash # function not queried by tag shouldn't have status assert function_queried_by_tag["status"] is not None assert function_queried_by_hash_key["status"] is None
def test_store_function_default_to_latest(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla"} function_name_1 = "function_name_1" function_hash_key = db.store_function(db_session, function_1, function_name_1) assert function_hash_key is not None function_queried_without_tag = db.get_function(db_session, function_name_1) function_queried_without_tag_hash = function_queried_without_tag[ "metadata"]["hash"] assert function_hash_key == function_queried_without_tag_hash assert function_queried_without_tag["metadata"]["tag"] == "latest" function_queried_with_tag = db.get_function(db_session, function_name_1, tag="latest") function_queried_without_tag_hash = function_queried_with_tag["metadata"][ "hash"] assert function_queried_with_tag is not None assert function_queried_with_tag["metadata"]["tag"] == "latest" assert function_queried_without_tag_hash == function_queried_without_tag_hash
def test_get_function_by_tag(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla', 'status': {'bla': 'blabla'}} function_name_1 = 'function_name_1' function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=True) function_queried_by_tag = db.get_function(db_session, function_name_1, tag='latest') function_queried_by_hash_key = db.get_function(db_session, function_name_1, hash_key=function_hash_key) function_not_queried_by_tag_hash = function_queried_by_hash_key[ 'metadata']['hash'] assert function_hash_key == function_not_queried_by_tag_hash # function not queried by tag shouldn't have status assert function_queried_by_tag['status'] is not None assert function_queried_by_hash_key['status'] is None
def test_list_artifact_iter_parameter(db: DBInterface, db_session: Session): artifact_name_1 = "artifact_name_1" artifact_name_2 = "artifact_name_2" artifact_1 = _generate_artifact(artifact_name_1) artifact_2 = _generate_artifact(artifact_name_2) uid = "artifact_uid" # Use iters with multiple digits, to make sure filtering them via regex works test_iters = [0, 5, 9, 42, 219, 2102] for iter in test_iters: artifact_1["iter"] = artifact_2["iter"] = iter db.store_artifact(db_session, artifact_name_1, artifact_1, uid, iter) db.store_artifact(db_session, artifact_name_2, artifact_2, uid, iter) # No filter on iter. All are expected artifacts = db.list_artifacts(db_session) assert len(artifacts) == len(test_iters) * 2 # Look for the various iteration numbers. Note that 0 is a special case due to the DB structure for iter in test_iters: artifacts = db.list_artifacts(db_session, iter=iter) assert len(artifacts) == 2 for artifact in artifacts: assert artifact["iter"] == iter # Negative test artifacts = db.list_artifacts(db_session, iter=666) assert len(artifacts) == 0 # Iter filter and a name filter, make sure query composition works artifacts = db.list_artifacts(db_session, name=artifact_name_1, iter=2102) assert len(artifacts) == 1
def test_list_functions_filtering_unversioned_untagged(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla"} function_2 = {"bla": "blablablabla"} function_name_1 = "function_name_1" tag = "some_tag" db.store_function(db_session, function_1, function_name_1, versioned=False, tag=tag) tagged_function_hash_key = db.store_function(db_session, function_2, function_name_1, versioned=True, tag=tag) functions = db.list_functions(db_session, function_name_1) # First we stored to the tag without versioning (unversioned instance) then we stored to the tag with version # so the unversioned instance remained untagged, verifying we're not getting it assert len(functions) == 1 assert functions[0]["metadata"]["hash"] == tagged_function_hash_key
def test_list_artifacts_exact_name_match(db: DBInterface, db_session: Session): artifact_1_key = "pre_artifact_key_suffix" artifact_2_key = "pre-artifact-key-suffix" artifact_1_uid = "artifact_uid_1" artifact_2_uid = "artifact_uid_2" artifact_1_body = _generate_artifact(artifact_1_key, uid=artifact_1_uid) artifact_2_body = _generate_artifact(artifact_2_key, uid=artifact_2_uid) # Store each twice - once with no iter, and once with an iter db.store_artifact( db_session, artifact_1_key, artifact_1_body, artifact_1_uid, ) artifact_1_body["iter"] = 42 db.store_artifact( db_session, artifact_1_key, artifact_1_body, artifact_1_uid, iter=42, ) db.store_artifact( db_session, artifact_2_key, artifact_2_body, artifact_2_uid, ) artifact_2_body["iter"] = 42 db.store_artifact( db_session, artifact_2_key, artifact_2_body, artifact_2_uid, iter=42, ) def _list_and_assert_count(key, count, iter=None): results = db.list_artifacts(db_session, name=key, iter=iter) assert len(results) == count return results # Ensure fuzzy query works, and we have everything we need _list_and_assert_count("~key", count=4) # Do an exact match with underscores in the name - must escape the _ do it doesn't do a like query list_results = _list_and_assert_count(artifact_1_key, count=2) for artifact in list_results: assert artifact["metadata"]["name"] == artifact_1_key _list_and_assert_count("%key%", count=0) # Verify we don't get artifacts whose name is "%-suffix" due to the like query used in the DB _list_and_assert_count("suffix", count=0) # This should also be filtered, since the prefix is "pre" which is 3 chars. There's a known caveat if # prefix is 1 or 2 chars long. _list_and_assert_count("artifact-key-suffix", count=0) _list_and_assert_count(artifact_1_key, iter=42, count=1) _list_and_assert_count("~key", iter=42, count=2) _list_and_assert_count("~key", iter=666, count=0)
def test_get_function_not_found(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla", "status": {"bla": "blabla"}} function_name_1 = "function_name_1" db.store_function(db_session, function_1, function_name_1, versioned=True) with pytest.raises(mlrun.errors.MLRunNotFoundError): db.get_function(db_session, function_name_1, tag="inexistent_tag") with pytest.raises(mlrun.errors.MLRunNotFoundError): db.get_function(db_session, function_name_1, hash_key="inexistent_hash_key")
def test_data_migration_fill_project_state( db: DBInterface, db_session: sqlalchemy.orm.Session, ): for i in range(10): project_name = f"project-name-{i}" _generate_and_insert_pre_060_record(db_session, project_name) projects = db.list_projects(db_session) for project in projects.projects: # getting default value from the schema assert project.spec.desired_state == mlrun.api.schemas.ProjectState.online assert project.status.state is None mlrun.api.initial_data._fill_project_state(db, db_session) projects = db.list_projects(db_session) for project in projects.projects: assert project.spec.desired_state == mlrun.api.schemas.ProjectState.online assert project.status.state == project.spec.desired_state # verify not storing for no reason db.store_project = unittest.mock.Mock() mlrun.api.initial_data._fill_project_state(db, db_session) assert db.store_project.call_count == 0
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 _ensure_runtime_resource_run_logs_collected(self, db: DBInterface, db_session: Session, runtime_resource: Dict): project, uid = self._resolve_runtime_resource_run(runtime_resource) # if cannot resolve related run, assume collected if not uid: return # import here to avoid circular imports import mlrun.api.crud as crud log_file_exists = crud.Logs.log_file_exists(project, uid) store_log = False if not log_file_exists: store_log = True else: log_mtime = crud.Logs.get_log_mtime(project, uid) log_mtime_datetime = datetime.fromtimestamp( log_mtime, timezone.utc) now = datetime.now(timezone.utc) run = db.read_run(db_session, uid, project) last_update_str = run.get('status', {}).get('last_update', now) last_update = datetime.fromisoformat(last_update_str) # this function is used to verify that logs collected from runtime resources before deleting them # here we're using the knowledge that the function is called only after a it was verified that the runtime # resource run is not in transient state, so we're assuming the run's last update is the last one, so if the # log file was modified after it, we're considering it as all logs collected if log_mtime_datetime < last_update: store_log = True if store_log: logger.debug(f'Storing runtime resource log before deletion') logs_from_k8s, _ = crud.Logs.get_log(db_session, project, uid, source=LogSources.K8S) crud.Logs.store_log(logs_from_k8s, project, uid, append=False)
def test_store_function_not_versioned(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla"} function_name_1 = "function_name_1" function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=False) function_result_1 = db.get_function(db_session, function_name_1) assert function_result_1 is not None assert function_result_1["metadata"]["tag"] == "latest" # not versioned so not queryable by hash key with pytest.raises(mlrun.errors.MLRunNotFoundError): db.get_function(db_session, function_name_1, hash_key=function_hash_key) function_2 = {"bla": "blabla", "bla2": "blabla2"} db.store_function(db_session, function_2, function_name_1, versioned=False) functions = db.list_functions(db_session, function_name_1) # Verifying only the latest version was saved assert len(functions) == 1
def test_store_function_not_versioned(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla'} function_name_1 = 'function_name_1' function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=False) function_result_1 = db.get_function(db_session, function_name_1) assert function_result_1 is not None assert function_result_1['metadata']['tag'] == 'latest' # not versioned so not queryable by hash key function_result_2 = db.get_function(db_session, function_name_1, hash_key=function_hash_key) assert function_result_2 is None function_2 = {'bla': 'blabla', 'bla2': 'blabla2'} db.store_function(db_session, function_2, function_name_1, versioned=False) functions = db.list_functions(db_session, function_name_1) # Verifying only the latest version was saved assert len(functions) == 1
def test_store_function_versioned(db: DBInterface, db_session: Session): function_1 = {'bla': 'blabla'} function_name_1 = 'function_name_1' function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=True) function_queried_without_hash_key = db.get_function( db_session, function_name_1) assert function_queried_without_hash_key is not None assert function_queried_without_hash_key['metadata']['tag'] == 'latest' # Verifying versioned function is queryable by hash_key function_queried_with_hash_key = db.get_function( db_session, function_name_1, hash_key=function_hash_key) function_queried_with_hash_key_hash = function_queried_with_hash_key[ 'metadata']['hash'] assert function_queried_with_hash_key is not None assert function_queried_with_hash_key['metadata']['tag'] == '' assert function_queried_with_hash_key_hash == function_hash_key function_2 = {'bla': 'blabla', 'bla2': 'blabla2'} function_name_1 = 'function_name_1' db.store_function(db_session, function_2, function_name_1, versioned=True) functions = db.list_functions(db_session, function_name_1) # Verifying both versions of the functions were saved assert len(functions) == 2 tagged_count = 0 for function in functions: if function['metadata']['tag'] == 'latest': tagged_count += 1 # but only one was tagged assert tagged_count == 1
def test_store_function_versioned(db: DBInterface, db_session: Session): function_1 = {"bla": "blabla"} function_name_1 = "function_name_1" function_hash_key = db.store_function(db_session, function_1, function_name_1, versioned=True) function_queried_without_hash_key = db.get_function( db_session, function_name_1) assert function_queried_without_hash_key is not None assert function_queried_without_hash_key["metadata"]["tag"] == "latest" # Verifying versioned function is queryable by hash_key function_queried_with_hash_key = db.get_function( db_session, function_name_1, hash_key=function_hash_key) function_queried_with_hash_key_hash = function_queried_with_hash_key[ "metadata"]["hash"] assert function_queried_with_hash_key is not None assert function_queried_with_hash_key["metadata"]["tag"] == "" assert function_queried_with_hash_key_hash == function_hash_key function_2 = {"bla": "blabla", "bla2": "blabla2"} function_name_1 = "function_name_1" db.store_function(db_session, function_2, function_name_1, versioned=True) functions = db.list_functions(db_session, function_name_1) # Verifying both versions of the functions were saved assert len(functions) == 2 tagged_count = 0 for function in functions: if function["metadata"]["tag"] == "latest": tagged_count += 1 # but only one was tagged assert tagged_count == 1
def test_list_artifact_category_filter(db: DBInterface, db_session: Session): artifact_name_1 = "artifact_name_1" artifact_kind_1 = ChartArtifact.kind artifact_name_2 = "artifact_name_2" artifact_kind_2 = PlotArtifact.kind artifact_name_3 = "artifact_name_3" artifact_kind_3 = ModelArtifact.kind artifact_name_4 = "artifact_name_4" artifact_kind_4 = DatasetArtifact.kind artifact_1 = _generate_artifact(artifact_name_1, kind=artifact_kind_1) artifact_2 = _generate_artifact(artifact_name_2, kind=artifact_kind_2) artifact_3 = _generate_artifact(artifact_name_3, kind=artifact_kind_3) artifact_4 = _generate_artifact(artifact_name_4, kind=artifact_kind_4) uid = "artifact_uid" db.store_artifact( db_session, artifact_name_1, artifact_1, uid, ) db.store_artifact( db_session, artifact_name_2, artifact_2, uid, ) db.store_artifact( db_session, artifact_name_3, artifact_3, uid, ) db.store_artifact( db_session, artifact_name_4, artifact_4, uid, ) artifacts = db.list_artifacts(db_session) assert len(artifacts) == 4 artifacts = db.list_artifacts(db_session, category=schemas.ArtifactCategories.model) assert len(artifacts) == 1 assert artifacts[0]["metadata"]["name"] == artifact_name_3 artifacts = db.list_artifacts(db_session, category=schemas.ArtifactCategories.dataset) assert len(artifacts) == 1 assert artifacts[0]["metadata"]["name"] == artifact_name_4 artifacts = db.list_artifacts(db_session, category=schemas.ArtifactCategories.other) assert len(artifacts) == 2 assert artifacts[0]["metadata"]["name"] == artifact_name_1 assert artifacts[1]["metadata"]["name"] == artifact_name_2