def test_upgrade_graph(self):
data = {'nodes': 'abc', 'edges': 'def', 'task_info': 'ghi'}
t1 = TaskInst()
t2 = TaskInst()
t1.save_graph = MagicMock()
t2.get_graph = MagicMock(return_value=data)
t1.upgrade_graph(t2)
t2.get_graph.assert_called_once()
t1.save_graph.assert_called_once_with('abc', 'def', 'ghi')
def test_save_graph(self, mock_assert_start_end, mock_get_sid_edge_sets,
mock_set_diff, mock_get_node_edge_map,
mock_remove_edges, mock_remove_nodes,
mock_change_edges, mock_change_nodes, mock_add_edges,
mock_add_nodes, mock_preprocess_edges,
mock_get_user_list):
task = TaskInst(name='task').save()
task.save_graph('my nodes', 'my edges')
mock_assert_start_end.assert_called_once_with('my nodes')
self.assertEqual(2, mock_get_sid_edge_sets.call_count)
self.assertEqual(2, mock_set_diff.call_count)
mock_get_node_edge_map.assert_called_once_with('my nodes',
'processed_edges')
mock_remove_edges.assert_called_once()
mock_remove_nodes.assert_called_once()
mock_change_edges.assert_called_once()
mock_change_nodes.assert_called_once()
mock_add_edges.assert_called_once()
mock_add_nodes.assert_called_once()