From 15344527aa8556798776aa0b3565af36e6c3f3a3 Mon Sep 17 00:00:00 2001
From: maybe-sybr <58414429+maybe-sybr@users.noreply.github.com>
Date: Mon, 27 Jul 2020 11:08:19 +1000
Subject: [PATCH] fix: Respect name of `@Custom*` decorated defs

---
 stix2/custom.py               | 8 ++++++++
 stix2/test/v20/test_custom.py | 2 +-
 stix2/test/v21/test_custom.py | 2 +-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/stix2/custom.py b/stix2/custom.py
index f749b04..08574ef 100644
--- a/stix2/custom.py
+++ b/stix2/custom.py
@@ -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
diff --git a/stix2/test/v20/test_custom.py b/stix2/test/v20/test_custom.py
index 33be3e2..70835c1 100644
--- a/stix2/test/v20/test_custom.py
+++ b/stix2/test/v20/test_custom.py
@@ -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)
diff --git a/stix2/test/v21/test_custom.py b/stix2/test/v21/test_custom.py
index 146abcd..ea6d3a8 100644
--- a/stix2/test/v21/test_custom.py
+++ b/stix2/test/v21/test_custom.py
@@ -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)