Merge pull request #3 from CybercentreCanada/submission_classificatio…

…n_propagation Submission classification propagation
CybercentreCanada · May 5, 2020 · 75bee07 · 75bee07
2 parents 8f8a185 + 24d8ddb
commit 75bee07
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/assemblyline_ui/api/v4/submission.py b/assemblyline_ui/api/v4/submission.py
@@ -375,10 +375,10 @@ def recursive_flatten_tree(tree):
         data.update(get_results(res_keys))
         data.update(get_errors(err_keys))
 
-        for r in data['results']:
+        for r in data['results'].values():
             data['classification'] = Classification.max_classification(data['classification'], r['classification'])
 
-        for f in data['file_infos']:
+        for f in data['file_infos'].values():
             data['classification'] = Classification.max_classification(data['classification'], f['classification'])
 
         return make_api_response(data)

diff --git a/test/test_submission.py b/test/test_submission.py
@@ -122,7 +122,11 @@ def test_get_submission_tree(datastore, login_session):
     submission = random.choice(datastore.submission.search("id:*", rows=NUM_SUBMISSIONS, as_obj=False)['items'])
     resp = get_api_data(session, f"{host}/api/v4/submission/tree/{submission['sid']}/")
     assert isinstance(resp, dict)
-    for k in resp:
+    assert "classification" in resp
+    assert "filtered" in resp
+    assert "tree" in resp
+
+    for k in resp['tree']:
         assert len(k) == 64