From 47b11453fa363b74d0e2213065e8c87d16a7ac0d Mon Sep 17 00:00:00 2001
From: Chris Lenk <clenk@mitre.org>
Date: Tue, 24 Oct 2017 12:43:30 -0400
Subject: [PATCH] Increase test coverage for memory datastore

---
 stix2/sources/__init__.py |  4 ++--
 stix2/sources/memory.py   |  4 ++--
 stix2/test/test_memory.py | 35 ++++++++++++++++++++++++++++++++++-
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/stix2/sources/__init__.py b/stix2/sources/__init__.py
index 47c7573..a142e88 100644
--- a/stix2/sources/__init__.py
+++ b/stix2/sources/__init__.py
@@ -77,7 +77,7 @@ class DataStore(object):
         """
         return self.source.all_versions(stix_id, allow_custom=allow_custom)
 
-    def query(self, query, allow_custom=False):
+    def query(self, query=None, allow_custom=False):
         """Retrieve STIX objects matching a set of filters.
 
         Implement: Specific data source API calls, processing,
@@ -191,7 +191,7 @@ class DataSource(object):
         """
         raise NotImplementedError()
 
-    def query(self, query, _composite_filters=None, allow_custom=False):
+    def query(self, query=None, _composite_filters=None, allow_custom=False):
         """
         Implement:Implement the specific data source API calls, processing,
         functionality required for retrieving query from the data source
diff --git a/stix2/sources/memory.py b/stix2/sources/memory.py
index 0f10cf9..0846f9b 100644
--- a/stix2/sources/memory.py
+++ b/stix2/sources/memory.py
@@ -65,7 +65,7 @@ def _add(store, stix_data=None, allow_custom=False):
             _add(store, stix_obj, allow_custom=allow_custom)
 
     else:
-        raise TypeError("stix_data must be as STIX object(or list of),json formatted STIX (or list of), or a json formatted STIX bundle")
+        raise TypeError("stix_data must be a STIX object (or list of), JSON formatted STIX (or list of), or a JSON formatted STIX bundle")
 
 
 class MemoryStore(DataStore):
@@ -283,7 +283,7 @@ class MemorySource(DataSource):
             if not isinstance(query, list):
                 # make sure dont make set from a Filter object,
                 # need to make a set from a list of Filter objects (even if just one Filter)
-                query = list(query)
+                query = [query]
             query = set(query)
 
         # combine all query filters
diff --git a/stix2/test/test_memory.py b/stix2/test/test_memory.py
index e8ba56b..0603bf7 100644
--- a/stix2/test/test_memory.py
+++ b/stix2/test/test_memory.py
@@ -138,11 +138,34 @@ def test_memory_store_all_versions(mem_store):
 
 def test_memory_store_query(mem_store):
     query = [Filter('type', '=', 'malware')]
-
     resp = mem_store.query(query)
     assert len(resp) == 0
 
 
+def test_memory_store_query_single_filter(mem_store):
+    query = Filter('id', '=', 'indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f')
+    resp = mem_store.query(query)
+    assert len(resp) == 1
+
+
+def test_memory_store_query_empty_query(mem_store):
+    resp = mem_store.query()
+    # sort since returned in random order
+    resp = sorted(resp, key=lambda k: k['id'])
+    assert len(resp) == 2
+    assert resp[0]['id'] == 'indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f'
+    assert resp[0]['modified'] == '2017-01-27T13:49:53.935Z'
+    assert resp[1]['id'] == 'indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f'
+    assert resp[1]['modified'] == '2017-01-27T13:49:53.936Z'
+
+
+def test_memory_store_query_multiple_filters(mem_store):
+    mem_store.source.filters.add(Filter('type', '=', 'indicator'))
+    query = Filter('id', '=', 'indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f')
+    resp = mem_store.query(query)
+    assert len(resp) == 1
+
+
 def test_memory_store_add_stix_object_str(mem_store):
     # add stix object string
     camp_id = "campaign--111111b6-1112-4fb0-111b-b111107ca70a"
@@ -194,6 +217,16 @@ def test_memory_store_add_stix_bundle_str(mem_store):
     assert camp_alias in camp_r["aliases"]
 
 
+def test_memory_store_add_invalid_object(mem_store):
+    ind = ('indicator', IND1)  # tuple isn't valid
+    with pytest.raises(TypeError) as excinfo:
+        mem_store.add(ind)
+    assert 'stix_data must be' in str(excinfo.value)
+    assert 'a STIX object' in str(excinfo.value)
+    assert 'JSON formatted STIX' in str(excinfo.value)
+    assert 'JSON formatted STIX bundle' in str(excinfo.value)
+
+
 def test_memory_store_object_with_custom_property(mem_store):
     camp = Campaign(name="Scipio Africanus",
                     objective="Defeat the Carthaginians",