:fix: from statement

linkdd · Aug 31, 2016 · a22b34f · a22b34f
1 parent a74aeee
commit a22b34f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 16 deletions.
diff --git a/link/graph/dsl/walker/test/walkthrough.py b/link/graph/dsl/walker/test/walkthrough.py
@@ -54,6 +54,20 @@ def mapper_emit(key, item):
         return result
 
     def setUp(self):
+        self.graphmgr = MagicMock()
+        self.graphmgr.mapreduce.side_effect = self._mapreduce
+        self.walk = Walkthrough(self.graphmgr)
+
+        patcher = patch('link.graph.dsl.walker.walkthrough.getfeature')
+        self.getfeature = patcher.start()
+        self.getfeature.side_effect = self._getfeature
+        self.addCleanup(patcher.stop)
+
+    def tearDown(self):
+        self.expected_nodes = {}
+        self.expected_rels = {}
+
+    def test_walkthrough(self):
         self.expected_nodes = {
             'fromfilter': [
                 {'_id': 'buzz'}
@@ -77,16 +91,6 @@ def setUp(self):
             ]
         }
 
-        self.graphmgr = MagicMock()
-        self.graphmgr.mapreduce.side_effect = self._mapreduce
-        self.walk = Walkthrough(self.graphmgr)
-
-        patcher = patch('link.graph.dsl.walker.walkthrough.getfeature')
-        self.getfeature = patcher.start()
-        self.getfeature.side_effect = self._getfeature
-        self.addCleanup(patcher.stop)
-
-    def test_walkthrough(self):
         from_ = MagicMock()
         from_.set_ = 'NODES'
         from_.alias = 'elt0'
@@ -117,6 +121,21 @@ def test_walkthrough(self):
 
         result = self.walk([statement])
 
+        self.assertIn('elt0', result)
+        self.assertIn('elt1', result)
+        self.assertIn('rel0', result)
+
+        self.assertEqual(result['elt0']['type'], 'nodes')
+
+        self.assertEqual(len(result['elt0']['dataset']), 1)
+        self.assertIn({'_id': 'buzz'}, result['elt0']['dataset'])
+
+        self.assertEqual(len(result['elt1']['dataset']), 1)
+        self.assertIn({'_id': 'sarge'}, result['elt1']['dataset'])
+
+        self.assertEqual(len(result['rel0']['dataset']), 1)
+        self.assertIn({'_id': 'andy'}, result['rel0']['dataset'])
+
 
 if __name__ == '__main__':
     main()
diff --git a/link/graph/dsl/walker/walkthrough.py b/link/graph/dsl/walker/walkthrough.py
@@ -41,8 +41,6 @@ def __call__(self, statements):
                 nodes = aliased_sets[last_alias]['dataset']
 
                 for stmt in path.through:
-                    local_nodes = nodes
-
                     if stmt.alias is not None and stmt.alias in aliased_sets:
                         rels = aliased_sets[stmt.alias]['dataset']
 
@@ -59,14 +57,12 @@ def __call__(self, statements):
                         }
 
                     if rels:
-                        local_nodes = self.walk_nodes(
-                            local_nodes,
+                        nodes = self.walk_nodes(
+                            nodes,
                             rels,
                             stmt.wmode
                         )
 
-                        nodes += local_nodes
-
                 for to in path.to:
                     aliased_sets[to.alias] = {
                         'type': 'nodes',