From e91b71f3009788f69c7303e22aca6a852c4cacf0 Mon Sep 17 00:00:00 2001
From: Chris Lenk <clenk@mitre.org>
Date: Wed, 14 Mar 2018 12:47:28 -0400
Subject: [PATCH] Test adding a data source to the workbench

---
 stix2/test/test_workbench.py | 24 +++++++++++++++++++++---
 stix2/workbench.py           |  6 ++++--
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/stix2/test/test_workbench.py b/stix2/test/test_workbench.py
index 6a33f11..b84b529 100644
--- a/stix2/test/test_workbench.py
+++ b/stix2/test/test_workbench.py
@@ -1,10 +1,12 @@
+import os
+
 import stix2
 from stix2.workbench import (AttackPattern, Campaign, CourseOfAction, Identity,
                              Indicator, IntrusionSet, Malware, ObservedData,
                              Report, ThreatActor, Tool, Vulnerability, add,
-                             all_versions, attack_patterns, campaigns,
-                             courses_of_action, create, get, identities,
-                             indicators, intrusion_sets, malware,
+                             add_data_source, all_versions, attack_patterns,
+                             campaigns, courses_of_action, create, get,
+                             identities, indicators, intrusion_sets, malware,
                              observed_data, query, reports, threat_actors,
                              tools, vulnerabilities)
 
@@ -171,3 +173,19 @@ def test_workbench_related():
     assert any(x['id'] == CAMPAIGN_ID for x in resp)
     assert any(x['id'] == INDICATOR_ID for x in resp)
     assert any(x['id'] == IDENTITY_ID for x in resp)
+
+    resp = get(MALWARE_ID).related(relationship_type='indicates')
+    assert len(resp) == 1
+
+
+def test_add_data_source():
+    fs_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "stix2_data")
+    fs = stix2.FileSystemSource(fs_path)
+    add_data_source(fs)
+
+    resp = tools()
+    assert len(resp) == 3
+    resp_ids = [tool.id for tool in resp]
+    assert TOOL_ID in resp_ids
+    assert 'tool--03342581-f790-4f03-ba41-e82e67392e23' in resp_ids
+    assert 'tool--242f3da3-4425-4d11-8f5c-b842886da966' in resp_ids
diff --git a/stix2/workbench.py b/stix2/workbench.py
index f5f2d32..60d1165 100644
--- a/stix2/workbench.py
+++ b/stix2/workbench.py
@@ -13,10 +13,11 @@ from . import Report as _Report
 from . import ThreatActor as _ThreatActor
 from . import Tool as _Tool
 from . import Vulnerability as _Vulnerability
+from .datastore.filters import Filter
+from .datastore.memory import MemoryStore
 from .environment import Environment
-from .sources.filters import Filter
-from .sources.memory import MemoryStore
 
+# Use an implicit MemoryStore
 _environ = Environment(store=MemoryStore())
 
 create = _environ.create
@@ -31,6 +32,7 @@ add_filters = _environ.add_filters
 add_filter = _environ.add_filter
 parse = _environ.parse
 add_data_source = _environ.source.add_data_source
+add_data_sources = _environ.source.add_data_sources
 
 
 # Wrap SDOs with helper functions