def test_entity_class_contains_labels():
entity = EntityV2(
"my-entity",
description="My entity",
value_type=ValueType.STRING,
labels={
"key1": "val1",
"key2": "val2"
},
)
assert "key1" in entity.labels.keys() and entity.labels["key1"] == "val1"
assert "key2" in entity.labels.keys() and entity.labels["key2"] == "val2"
def test_entity_import_export_yaml(self):
test_entity = EntityV2(
name="car_driver_entity",
description="Driver entity for car rides",
value_type=ValueType.STRING,
labels={"team": "matchmaking"},
)
# Create a string YAML representation of the entity
string_yaml = test_entity.to_yaml()
# Create a new entity object from the YAML string
actual_entity_from_string = EntityV2.from_yaml(string_yaml)
# Ensure equality is upheld to original entity
assert test_entity == actual_entity_from_string
def test_entity_without_labels_empty_dict():
entity = EntityV2("my-entity",
description="My entity",
value_type=ValueType.STRING)
assert entity.labels == dict()
assert len(entity.labels) == 0