Skip to content

Commit

Permalink
Update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
vincenzocaputo committed Apr 2, 2024
1 parent 17acbad commit a10bfbd
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 66 deletions.
2 changes: 1 addition & 1 deletion tests/bundle1_graph.json

Large diffs are not rendered by default.

7 changes: 4 additions & 3 deletions tests/graph.html
Original file line number Diff line number Diff line change
Expand Up @@ -159,10 +159,11 @@ <h1></h1>
const selectedNode = network.getSelectedNodes()[0];
if (allNodes[selectedNode]) {
document.querySelector("#code_section").style.display="block";
document.querySelector("#code_section pre").textContent =
JSON.stringify(JSON.parse(allNodes[selectedNode].stix), null, 2);
document.querySelector("#code_section pre").textContent =
JSON.stringify(JSON.parse(
allNodes[selectedNode].stix), null, 2);
}
}
}
})
</script></body>
</html>
131 changes: 69 additions & 62 deletions tests/test_pystixview.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,17 +13,8 @@

class TestPySTIXView(unittest.TestCase):

def test_empty_network(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
with open('tests/empty_graph.json', 'r') as fd:
self.assertEqual(self.graph.to_json(), fd.read())

def test_empty_network_fail(self):
with self.assertRaises(ValueError):
self.graph = PySTIXView("100%", "100%", notebook=False, style="test")

def test_add_node_observabel(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
def test_add_node_observabe(self):
self.graph = PySTIXView(notebook=False)

stix_obs = """{
"type": "domain-name",
Expand All @@ -33,27 +24,27 @@ def test_add_node_observabel(self):
}
"""
self.assertTrue(self.graph.add_node(stix_obs))
network = json.loads(self.graph.to_json())
self.assertEqual(network[0][0]['label'], "test-domain.local")
graph = json.loads(self.graph.to_json())
self.assertEqual(next(iter(graph['nodes'].values()))['stix_type'], "observable")

def test_bundle_string(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle1.json', 'r') as fd:
self.graph.add_bundle(fd.read())

with open('tests/bundle1_graph.json', 'r') as fd:
self.assertEqual(self.graph.to_json(), fd.read())

def test_bundle_dict(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle1.json', 'r') as fd:
self.graph.add_bundle(json.loads(fd.read()))

with open('tests/bundle1_graph.json', 'r') as fd:
self.assertEqual(self.graph.to_json(), fd.read())

def test_bundle_object(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle1.json', 'r') as fd:
bundle = parse(fd.read())
self.graph.add_bundle(bundle)
Expand All @@ -62,30 +53,30 @@ def test_bundle_object(self):
self.assertEqual(self.graph.to_json(), fd.read())

def test_bundle_fail(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with self.assertRaises(TypeError):
self.graph.add_bundle(123)

def test_bundle_granular_markings(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle3.json', 'r') as fd:
bundle = parse(fd.read())

self.graph.add_bundle(bundle)
network = json.loads(self.graph.to_json())
self.assertTrue('TLP:RED' in [ x['label'] for x in network[0] ])
graph = json.loads(self.graph.to_json())
self.assertTrue('TLP:RED' in [ x['name'] for x in list(graph['nodes'].values()) ])

def test_bundle_object_refs(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle4.json', 'r') as fd:
bundle = parse(fd.read())

self.graph.add_bundle(bundle)
network = json.loads(self.graph.to_json())
self.assertEqual('malware--02524b02-07a2-4b38-a449-e1be2852ed97', network[1][0]['to'])
graph = json.loads(self.graph.to_json())
self.assertEqual('malware--02524b02-07a2-4b38-a449-e1be2852ed97', [ x['to'] for x in graph['edges'] ][0])

def test_add_relationship_str(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle2.json', 'r') as fd:
bundle = parse(fd.read())
self.graph.add_bundle(bundle)
Expand All @@ -100,10 +91,12 @@ def test_add_relationship_str(self):
"target_ref": "malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4"
}"""
self.graph.add_relationship(relationship)
self.assertTrue( "\"from\": \"threat-actor--9a8a0d25-7636-429b-a99e-b2a73cd0f11f\", \"to\": \"malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4\"" in self.graph.to_json() )
graph = json.loads(self.graph.to_json())
self.assertEqual('threat-actor--9a8a0d25-7636-429b-a99e-b2a73cd0f11f', [ x['from'] for x in graph['edges'] ][0])
self.assertEqual('malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4', [ x['to'] for x in graph['edges'] ][0])

def test_add_relationship_dict(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle2.json', 'r') as fd:
bundle = parse(fd.read())
self.graph.add_bundle(bundle)
Expand All @@ -118,10 +111,12 @@ def test_add_relationship_dict(self):
"target_ref": "malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4"
}"""
self.graph.add_relationship(json.loads(relationship))
self.assertTrue( "\"from\": \"threat-actor--9a8a0d25-7636-429b-a99e-b2a73cd0f11f\", \"to\": \"malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4\"" in self.graph.to_json() )
graph = json.loads(self.graph.to_json())
self.assertEqual('threat-actor--9a8a0d25-7636-429b-a99e-b2a73cd0f11f', [ x['from'] for x in graph['edges'] ][0])
self.assertEqual('malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4', [ x['to'] for x in graph['edges'] ][0])

def test_add_relationship_stix(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with open('tests/stix_bundle2.json', 'r') as fd:
bundle = parse(fd.read())
self.graph.add_bundle(bundle)
Expand All @@ -131,20 +126,22 @@ def test_add_relationship_stix(self):
target_ref="malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4")

self.graph.add_relationship(relationship)
self.assertTrue( "\"from\": \"threat-actor--9a8a0d25-7636-429b-a99e-b2a73cd0f11f\", \"to\": \"malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4\"" in self.graph.to_json() )
graph = json.loads(self.graph.to_json())
self.assertEqual('threat-actor--9a8a0d25-7636-429b-a99e-b2a73cd0f11f', [ x['from'] for x in graph['edges'] ][0])
self.assertEqual('malware--d1c612bc-146f-4b65-b7b0-9a54a14150a4', [ x['to'] for x in graph['edges'] ][0])

def test_add_relationship_fail(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

with self.assertRaises(TypeError):
self.graph.add_relationship(123)

def test_custom_type_fail(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
self.assertRaises(TypeError, self.graph.add_custom_stix_type("x-test"))

def test_custom_type_icon(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
self.graph.add_custom_stix_type("x-test", node_icon="tests/test-icon.png")
self.graph.add_node("""
{
Expand All @@ -157,12 +154,13 @@ def test_custom_type_icon(self):
"description": "Test"
}
""")
graph_html = self.graph.show_graph(width="100%", height="100%")
with open('tests/test-icon.png', 'rb') as fd:
img_b64 = f"data:image/png;base64,{base64.b64encode(fd.read()).decode('utf-8')}"
self.assertTrue( img_b64 in self.graph.to_json() )
self.assertTrue( img_b64 in graph_html )

def test_custom_type_icon_url(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

url = "https://uxwing.com/wp-content/themes/uxwing/download/web-app-development/bug-icon.png"
self.graph.add_custom_stix_type("x-test", node_icon=url)
Expand All @@ -177,10 +175,11 @@ def test_custom_type_icon_url(self):
"description": "Test"
}
""")
self.assertTrue( url in self.graph.to_json() )
graph_html = self.graph.show_graph(width="100%", height="100%")
self.assertTrue( url in graph_html )

def test_custom_type_icon_fail(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

with self.assertRaises(TypeError):
self.graph.add_custom_stix_type("x-test", node_icon=123)
Expand All @@ -189,15 +188,15 @@ def test_custom_type_icon_fail(self):
self.graph.add_custom_stix_type("x-test-2", node_color="test")

def test_custom_type_fail(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

with self.assertRaises(Exception):
self.graph.add_custom_stix_type("x-test", node_color="#FF0000")
self.graph.add_custom_stix_type("x-test", node_color="#FF0000")


def test_custom_type_icon_color(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

self.graph.add_custom_stix_type("x-test", node_color="#ABCDEF")
self.graph.add_node("""
Expand All @@ -211,50 +210,58 @@ def test_custom_type_icon_color(self):
"description": "Test"
}
""")
self.assertTrue( "#ABCDEF" in self.graph.to_json() )
graph_html = self.graph.show_graph(width="100%", height="100%")
self.assertTrue( "#ABCDEF" in graph_html )

def test_node_type(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

self.graph.add_node("""
{
"type": "x-test",
"spec_version": "2.1",
"id": "x-test--45d71d3d-52bd-4300-815a-78c434d4d50c",
"created": "2024-03-03T15:11:57",
"modified": "2024-03-03T15:11:57",
"name": "Test",
"description": "Test"
}
""")
with self.assertWarns(Warning):
self.graph.add_node("""
{
"type": "x-test",
"spec_version": "2.1",
"id": "x-test--45d71d3d-52bd-4300-815a-78c434d4d50c",
"created": "2024-03-03T15:11:57",
"modified": "2024-03-03T15:11:57",
"name": "Test",
"description": "Test"
}
""")
self.graph.show_graph(width="100%", height="100%")

def test_add_node_fail(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)
with self.assertRaises(TypeError):
self.graph.add_node(123)


def test_show_graph(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

html = self.graph.show_graph(show_physics_buttons=True,
show_node_buttons=True,
show_edge_buttons=True)
html = self.graph.show_graph(width="100%",
height="100%",
show_physics_buttons=True,
show_node_buttons=True,
show_edge_buttons=True)
self.assertTrue( isinstance(html, str))

def test_show_graph_notebook(self):
self.graph = PySTIXView("100%", "100%", notebook=True)
self.graph = PySTIXView(notebook=True)

html = self.graph.show_graph(show_physics_buttons=True,
show_node_buttons=True,
show_edge_buttons=True)
html = self.graph.show_graph(width="100%",
height="100%",
show_physics_buttons=True,
show_node_buttons=True,
show_edge_buttons=True)
self.assertTrue( isinstance(html, HTML) )

def test_save_graph_notebook(self):
self.graph = PySTIXView("100%", "100%", notebook=False)
self.graph = PySTIXView(notebook=False)

self.graph.save_graph('tests/graph.html', show_physics_buttons=True,
self.graph.save_graph('tests/graph.html', height="100%",
width="100%",
show_physics_buttons=True,
show_node_buttons=True,
show_edge_buttons=True)

Expand Down

0 comments on commit a10bfbd

Please sign in to comment.