def test_element_set_at(sample_target_def: target.TargetDefinition): """Test element get method.""" element = Element(sample_target_def) metadata = target.Metadata( **{ 'title': 'My simple catalog', 'last-modified': datetime.now(), 'version': '0.0.0', 'oscal-version': '1.0.0-Milestone3' } ) parties: List[target.Party] = [] parties.append( target.Party(**{ 'uuid': 'ff47836c-877c-4007-bbf3-c9d9bd805000', 'name': 'TEST1', 'type': 'organization' }) ) parties.append( target.Party(**{ 'uuid': 'ee88836c-877c-4007-bbf3-c9d9bd805000', 'name': 'TEST2', 'type': 'organization' }) ) assert element.set_at(ElementPath('target-definition.metadata'), metadata).get_at(ElementPath('target-definition.metadata')) == metadata assert element.set_at(ElementPath('target-definition.metadata.parties'), parties).get_at(ElementPath('target-definition.metadata.parties')) == parties assert element.set_at(ElementPath('target-definition.metadata.parties.*'), parties).get_at(ElementPath('target-definition.metadata.parties')) == parties # unset assert element.set_at(ElementPath('target-definition.metadata.parties'), None).get_at(ElementPath('target-definition.metadata.parties')) is None # string element path assert element.set_at('target-definition.metadata.parties', parties).get_at(ElementPath('target-definition.metadata.parties')) == parties with pytest.raises(TrestleError): assert element.set_at(ElementPath('target-definition.metadata'), parties).get_at(ElementPath('target-definition.metadata.parties')) == parties # wildcard requires it to be an OscalBaseModel or list with pytest.raises(TrestleError): assert element.set_at(ElementPath('target-definition.metadata.parties.*'), 'INVALID') # invalid attribute with pytest.raises(TrestleError): assert element.set_at(ElementPath('target-definition.metadata.groups.*'), parties)
def add(element_path: ElementPath, parent_element: Element, include_optional: bool) -> None: """For a element_path, add a child model to the parent_element of a given parent_model. Args: element_path: element path of the item to create within the model parent_element: the parent element that will host the created element include_optional: whether to create optional attributes in the created element Notes: First we find the child model at the specified element path and instantiate it with default values. Then we check if there's already existing element at that path, in which case we append the child model to the existing list of dict. Then we set up an action plan to update the model (specified by file_path) in memory, create a file at the same location and write the file. We update the parent_element to prepare for next adds in the chain """ if '*' in element_path.get_full_path_parts(): raise err.TrestleError('trestle add does not support Wildcard element path.') # Get child model try: child_model = element_path.get_type(type(parent_element.get())) # Create child element with sample values child_object = gens.generate_sample_model(child_model, include_optional=include_optional) if parent_element.get_at(element_path) is not None: # The element already exists if type(parent_element.get_at(element_path)) is list: child_object = parent_element.get_at(element_path) + child_object elif type(parent_element.get_at(element_path)) is dict: child_object = {**parent_element.get_at(element_path), **child_object} else: raise err.TrestleError('Already exists and is not a list or dictionary.') except Exception as e: raise err.TrestleError(f'Bad element path. {str(e)}') update_action = UpdateAction( sub_element=child_object, dest_element=parent_element, sub_element_path=element_path ) parent_element = parent_element.set_at(element_path, child_object) return update_action, parent_element