Merge pull request #97 from oasis-open/memory-tests

Test MemoryStore saving/loading to/from file

Author: gtback

stix2/sources/memory.py

@@ -164,7 +164,7 @@ def save_to_file(self, file_path, allow_custom=False):
         if not os.path.exists(os.path.dirname(file_path)):
             os.makedirs(os.path.dirname(file_path))
         with open(file_path, "w") as f:
-            f.write(str(Bundle(self._data.values(), allow_custom=allow_custom)))
+            f.write(str(Bundle(list(self._data.values()), allow_custom=allow_custom)))
     save_to_file.__doc__ = MemoryStore.save_to_file.__doc__
 
 

stix2/test/test_memory.py

@@ -1,3 +1,6 @@
+import os
+import shutil
+
 import pytest
 
 from stix2 import (Bundle, Campaign, CustomObject, Filter, MemorySource,
@@ -166,6 +169,22 @@ def test_memory_store_query_multiple_filters(mem_store):
     assert len(resp) == 1
 
 
+def test_memory_store_save_load_file(mem_store):
+    filename = 'memory_test/mem_store.json'
+    mem_store.save_to_file(filename)
+    contents = open(os.path.abspath(filename)).read()
+
+    assert '"id": "indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f",' in contents
+    assert '"id": "indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f",' in contents
+
+    mem_store2 = MemoryStore()
+    mem_store2.load_from_file(filename)
+    assert mem_store2.get("indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f")
+    assert mem_store2.get("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")
+
+    shutil.rmtree(os.path.dirname(filename))
+
+
 def test_memory_store_add_stix_object_str(mem_store):
     # add stix object string
     camp_id = "campaign--111111b6-1112-4fb0-111b-b111107ca70a"