Improve error handling in relationships/related_to

stix2/sources/__init__.py

@@ -281,8 +281,11 @@ class DataSource(with_metaclass(ABCMeta)):
         filters = [Filter('type', '=', 'relationship')]
 
         try:
-            obj_id = obj.get('id', '')
-        except AttributeError:
+            obj_id = obj['id']
+        except KeyError:
+            raise ValueError("STIX object has no 'id' property")
+        except TypeError:
+            # Assume `obj` is an ID string
             obj_id = obj
 
         if relationship_type:
@@ -323,8 +326,11 @@ class DataSource(with_metaclass(ABCMeta)):
         rels = self.relationships(obj, relationship_type, source_only, target_only)
 
         try:
-            obj_id = obj.get('id', '')
-        except AttributeError:
+            obj_id = obj['id']
+        except KeyError:
+            raise ValueError("STIX object has no 'id' property")
+        except TypeError:
+            # Assume `obj` is an ID string
             obj_id = obj
 
         for r in rels:
@@ -533,8 +539,11 @@ class CompositeDataSource(DataSource):
         filters = [Filter('type', '=', 'relationship')]
 
         try:
-            obj_id = obj.get('id', '')
-        except AttributeError:
+            obj_id = obj['id']
+        except KeyError:
+            raise ValueError("STIX object has no 'id' property")
+        except TypeError:
+            # Assume `obj` is an ID string
             obj_id = obj
 
         if relationship_type:
@@ -583,8 +592,11 @@ class CompositeDataSource(DataSource):
             rels = ds.relationships(obj, relationship_type, source_only, target_only)
 
         try:
-            obj_id = obj.get('id', '')
-        except AttributeError:
+            obj_id = obj['id']
+        except KeyError:
+            raise ValueError("STIX object has no 'id' property")
+        except TypeError:
+            # Assume `obj` is an ID string
             obj_id = obj
 
         for ds in self.data_sources:

stix2/test/test_environment.py

@@ -258,6 +258,17 @@ def test_relationships(ds):
     assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp)
 
 
+def test_relationships_no_id(ds):
+    env = stix2.Environment(store=ds)
+    mal = {
+        "type": "malware",
+        "name": "some variant"
+    }
+    with pytest.raises(ValueError) as excinfo:
+        env.relationships(mal)
+    assert "object has no 'id' property" in str(excinfo.value)
+
+
 def test_relationships_by_type(ds):
     env = stix2.Environment(store=ds)
     mal = env.get(MALWARE_ID)
@@ -311,6 +322,17 @@ def test_related_to(ds):
     assert any(x['id'] == IDENTITY_ID for x in resp)
 
 
+def test_related_to_no_id(ds):
+    env = stix2.Environment(store=ds)
+    mal = {
+        "type": "malware",
+        "name": "some variant"
+    }
+    with pytest.raises(ValueError) as excinfo:
+        env.related_to(mal)
+    assert "object has no 'id' property" in str(excinfo.value)
+
+
 def test_related_to_by_source(ds):
     env = stix2.Environment(store=ds)
     resp = env.related_to(MALWARE_ID, source_only=True)