diff --git a/stix2/base.py b/stix2/base.py index 35360af..a44d879 100644 --- a/stix2/base.py +++ b/stix2/base.py @@ -23,13 +23,7 @@ class STIXJSONEncoder(json.JSONEncoder): def get_required_properties(properties): - for k, v in properties.items(): - if isinstance(v, dict): - if v.get('required'): - yield k - else: # This is a Property subclass - if v.required: - yield k + return (k for k, v in properties.items() if v.required) class _STIXBase(collections.Mapping): diff --git a/stix2/test/test_base.py b/stix2/test/test_base.py new file mode 100644 index 0000000..38a3e45 --- /dev/null +++ b/stix2/test/test_base.py @@ -0,0 +1,25 @@ +import datetime as dt +import json + +import pytest +import pytz + +from stix2.base import STIXJSONEncoder + + +def test_encode_json_datetime(): + now = dt.datetime(2017, 3, 22, 0, 0, 0, tzinfo=pytz.UTC) + test_dict = {'now': now} + + expected = '{"now": "2017-03-22T00:00:00Z"}' + assert json.dumps(test_dict, cls=STIXJSONEncoder) == expected + + +def test_encode_json_object(): + obj = object() + test_dict = {'obj': obj} + + with pytest.raises(TypeError) as excinfo: + json.dumps(test_dict, cls=STIXJSONEncoder) + + assert str(excinfo.value) == "Object of type 'object' is not JSON serializable"