From c4939368ae6be576d8209b21c6c8d3242bf8aa65 Mon Sep 17 00:00:00 2001
From: Matt Ball <mball@verily.com>
Date: Wed, 5 Jun 2019 15:12:50 -0600
Subject: [PATCH] update linker.py to be compatible with networkx-2.x

Also, update test_linker.py to run under Python 2 unit tests and add
test coverage for the modified functions.

linker.py should now be compatible with both networkx-1 and -2.

This addresses https://github.com/fishtown-analytics/dbt/issues/1496

This change was tested via unit tests, but was not tested via
integration tests.

As a follow-up change, core/setup.py should have its networkx dependency
updated to allow version 2.x.
---
 core/dbt/linker.py       |  4 ++--
 test/unit/test_linker.py | 23 ++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/core/dbt/linker.py b/core/dbt/linker.py
index 1a17c791af1..f7241bb68ff 100644
--- a/core/dbt/linker.py
+++ b/core/dbt/linker.py
@@ -136,7 +136,7 @@ def _find_new_additions(self):
 
         Callers must hold the lock.
         """
-        for node, in_degree in self.graph.in_degree_iter():
+        for node, in_degree in dict(self.graph.in_degree()).items():
             if not self._already_known(node) and in_degree == 0:
                 self.inner.put((self._scores[node], node))
                 self.queued.add(node)
@@ -268,5 +268,5 @@ def _updated_graph(graph, manifest):
         for key in GRAPH_SERIALIZE_BLACKLIST:
             if key in data:
                 del data[key]
-        graph.add_node(node_id, data)
+        graph.add_node(node_id, **data)
     return graph
diff --git a/test/unit/test_linker.py b/test/unit/test_linker.py
index 664df9455f6..6f76b90dd48 100644
--- a/test/unit/test_linker.py
+++ b/test/unit/test_linker.py
@@ -1,4 +1,6 @@
 import mock
+import os
+import tempfile
 import unittest
 
 import dbt.utils
@@ -6,7 +8,7 @@
 from dbt import linker
 try:
     from queue import Empty
-except KeyError:
+except ImportError:
     from Queue import Empty
 
 
@@ -37,6 +39,25 @@ def test_linker_add_node(self):
 
         self.assertEqual(len(actual_nodes), len(expected_nodes))
 
+    def test_linker_write_and_read_graph(self):
+        expected_nodes = ['A', 'B', 'C']
+        for node in expected_nodes:
+            self.linker.add_node(node)
+
+        manifest = _mock_manifest('ABC')
+        (fd, fname) = tempfile.mkstemp()
+        self.linker.write_graph(fname, manifest)
+
+        new_linker = linker.from_file(fname)
+        os.close(fd)
+        os.unlink(fname)
+
+        actual_nodes = new_linker.nodes()
+        for node in expected_nodes:
+            self.assertIn(node, actual_nodes)
+
+        self.assertEqual(len(actual_nodes), len(expected_nodes))
+
     def assert_would_join(self, queue):
         """test join() without timeout risk"""
         self.assertEqual(queue.inner.unfinished_tasks, 0)