From f615161110ce5444ec9bd0eb7e04b572a0d35e1b Mon Sep 17 00:00:00 2001
From: Michael Chisholm <chisholm@mitre.org>
Date: Tue, 6 Nov 2018 16:15:33 -0500
Subject: [PATCH] Added some tests for adding markings to sinks and stores.

---
 stix2/test/test_datastore_filesystem.py | 39 ++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/stix2/test/test_datastore_filesystem.py b/stix2/test/test_datastore_filesystem.py
index 44a12ed..0f0ca0a 100644
--- a/stix2/test/test_datastore_filesystem.py
+++ b/stix2/test/test_datastore_filesystem.py
@@ -10,7 +10,8 @@ import pytz
 
 from stix2 import (Bundle, Campaign, CustomObject, FileSystemSink,
                    FileSystemSource, FileSystemStore, Filter, Identity,
-                   Indicator, Malware, Relationship, parse, properties)
+                   Indicator, Malware, MarkingDefinition, Relationship,
+                   TLPMarking, parse, properties)
 from stix2.datastore.filesystem import (AuthSet, _find_search_optimizations,
                                         _get_matching_dir_entries,
                                         _timestamp2filename)
@@ -379,6 +380,22 @@ def test_filesystem_sink_add_objects_list(fs_sink, fs_source):
     os.remove(camp7filepath)
 
 
+def test_filesystem_sink_marking(fs_sink):
+    marking = MarkingDefinition(
+        definition_type="tlp",
+        definition=TLPMarking(tlp="green")
+    )
+
+    fs_sink.add(marking)
+    marking_filepath = os.path.join(
+        FS_PATH, "marking-definition", marking["id"] + ".json"
+    )
+
+    assert os.path.exists(marking_filepath)
+
+    os.remove(marking_filepath)
+
+
 def test_filesystem_store_get_stored_as_bundle(fs_store):
     coa = fs_store.get("course-of-action--95ddb356-7ba0-4bd9-a889-247262b8946f")
     assert coa.id == "course-of-action--95ddb356-7ba0-4bd9-a889-247262b8946f"
@@ -488,6 +505,26 @@ def test_filesystem_store_add_invalid_object(fs_store):
     assert 'JSON formatted STIX bundle' in str(excinfo.value)
 
 
+def test_filesystem_store_add_marking(fs_store):
+    marking = MarkingDefinition(
+        definition_type="tlp",
+        definition=TLPMarking(tlp="green")
+    )
+
+    fs_store.add(marking)
+    marking_filepath = os.path.join(
+        FS_PATH, "marking-definition", marking["id"] + ".json"
+    )
+
+    assert os.path.exists(marking_filepath)
+
+    marking_r = fs_store.get(marking["id"])
+    assert marking_r["id"] == marking["id"]
+    assert marking_r["definition"]["tlp"] == "green"
+
+    os.remove(marking_filepath)
+
+
 def test_filesystem_object_with_custom_property(fs_store):
     camp = Campaign(name="Scipio Africanus",
                     objective="Defeat the Carthaginians",