fix: Respect name of `@Custom*` decorated defs

pull/1/head
maybe-sybr 2020-07-27 11:08:19 +10:00
parent bcd3eb3e4b
commit 15344527aa
3 changed files with 10 additions and 2 deletions

View File

@ -35,6 +35,8 @@ def _custom_object_builder(cls, type, properties, version, base_class):
base_class.__init__(self, **kwargs)
_cls_init(cls, self, kwargs)
_CustomObject.__name__ = cls.__name__
_register_object(_CustomObject, version=version)
return _CustomObject
@ -51,6 +53,8 @@ def _custom_marking_builder(cls, type, properties, version, base_class):
base_class.__init__(self, **kwargs)
_cls_init(cls, self, kwargs)
_CustomMarking.__name__ = cls.__name__
_register_marking(_CustomMarking, version=version)
return _CustomMarking
@ -72,6 +76,8 @@ def _custom_observable_builder(cls, type, properties, version, base_class, id_co
base_class.__init__(self, **kwargs)
_cls_init(cls, self, kwargs)
_CustomObservable.__name__ = cls.__name__
_register_observable(_CustomObservable, version=version)
return _CustomObservable
@ -88,5 +94,7 @@ def _custom_extension_builder(cls, observable, type, properties, version, base_c
base_class.__init__(self, **kwargs)
_cls_init(cls, self, kwargs)
_CustomExtension.__name__ = cls.__name__
_register_observable_extension(observable, _CustomExtension, version=version)
return _CustomExtension

View File

@ -723,7 +723,7 @@ def test_custom_extension():
with pytest.raises(stix2.exceptions.MissingPropertiesError) as excinfo:
NewExtension(property2=42)
assert excinfo.value.properties == ['property1']
assert str(excinfo.value) == "No values for required properties for _CustomExtension: (property1)."
assert str(excinfo.value) == "No values for required properties for NewExtension: (property1)."
with pytest.raises(ValueError) as excinfo:
NewExtension(property1='something', property2=4)

View File

@ -920,7 +920,7 @@ def test_custom_extension():
with pytest.raises(stix2.exceptions.MissingPropertiesError) as excinfo:
NewExtension(property2=42)
assert excinfo.value.properties == ['property1']
assert str(excinfo.value) == "No values for required properties for _CustomExtension: (property1)."
assert str(excinfo.value) == "No values for required properties for NewExtension: (property1)."
with pytest.raises(ValueError) as excinfo:
NewExtension(property1='something', property2=4)