DB_row, update DBO_XXX iteration loops, test updates

3 files changed, 29 insertions, 12 deletions

diff --git a/src-py/db_controller.py b/src-py/db_controller.py
index 49ed3b6b..f4bfbb75 100644
--- a/src-py/db_controller.py
+++ b/src-py/db_controller.py
@@ -67,7 +67,8 @@ class DB_op(object):
         ret = []
         for _, _, r_set in self:
             for row in r_set:
-                ret.append(row)
+                for cloumn in row:
+                    ret.append(row)
         return ret
 
     def parse_multi_statement_response_data(self, data):
@@ -155,6 +156,7 @@ class DBO_add_node_set(DB_op):
         DB op: add node set
         
         @param node_map: node-type to node-set map
+        @return: set of new node DB ids
         """
         super(DBO_add_node_set, self).__init__()
         for q, q_param_set in db_util.gen_query_create_from_node_map(node_map):
@@ -164,8 +166,8 @@ class DBO_add_node_set(DB_op):
         id_set = []
         for _, _, r_set in self:
             for row in r_set:
-                nid = row  # [!] fragile
-                id_set.append(nid)
+                for clo in row:
+                    id_set.append(clo)
 
         log.debug('node-set added: ids: ' + str(id_set))
         return id_set
@@ -174,6 +176,7 @@ class DBO_add_link_set(DB_op):
     def __init__(self, link_map):
         """
         @param link_map: is a link-type to link-set map - see model.link
+        @return: set of new node DB ids
         """
         super(DBO_add_link_set, self).__init__()
         for q, q_params in db_util.gen_query_create_from_link_map(link_map):
@@ -181,10 +184,10 @@ class DBO_add_link_set(DB_op):
 
     def on_completion(self, data):
         id_set = []
-        for s_id, s, r_set in self:
+        for _, _, r_set in self:
             for row in r_set:
-                lid = row  # [!] fragile
-                id_set.append(lid)
+                for col_val in row:
+                    id_set.append(col_val)
 
         log.debug('link-set added: ids: ' + str(id_set))
         return id_set
diff --git a/src-py/neo4j_util.py b/src-py/neo4j_util.py
index 5169af0d..8d70ced9 100644
--- a/src-py/neo4j_util.py
+++ b/src-py/neo4j_util.py
@@ -7,6 +7,14 @@ import urllib2
 import model
 import string
 
+class DB_row(object):
+    def __init__(self, data):
+        self.data = data
+
+    def __iter__(self):
+        for column_val in self.data:
+            yield column_val
+
 class DB_result_set(object):
     def __init__(self, data):
         self.data = data
@@ -14,9 +22,9 @@ class DB_result_set(object):
     def __iter__(self):
         for db_row_dict in self.data['data']:
             # example: dict: {u'row': [{u'title': u'foo'}]}
-            assert None != db_row_dict['row'][0]
+            assert None != db_row_dict['row']
 
-            yield db_row_dict['row'][0]
+            yield DB_row(db_row_dict['row'])
 
 class Cypher_String_Formatter(string.Formatter):
     """
diff --git a/src-py/test_db_controller.py b/src-py/test_db_controller.py
index 66ecfd09..edfd4a22 100644
--- a/src-py/test_db_controller.py
+++ b/src-py/test_db_controller.py
@@ -135,13 +135,19 @@ class TestDBController(unittest.TestCase):
         n_set = self.db_ctl.exec_op(op)
         self.assertEqual(len(n_set), 1)
 
-    def test_node_DB_id_lifecycle(self):
+    def test_load_node_set_by_DB_id(self):
         """
         test node DB id life cycle
         """
-        id_set = self.db_ctl.exec_op(dbc.DBO_add_node_set({'Person': [{'name': 'John Doe', 'id': 'jdoe_00'},
-                                                                      {'name': 'John Doe', 'id': 'jdoe_01'}]}))
-        n_set = self.db_ctl.exec_op(dbc.DBO_load_node_set_by_DB_id(id_set))
+        
+        # create nodes, get DB ids
+        op = dbc.DBO_add_node_set({'T_test_load_node_set_by_DB_id': [{'name': 'John Doe'},
+                                                                     {'name': 'John Doe'}]})
+        id_set = self.db_ctl.exec_op(op)
+        
+        # match against DB ids
+        op = dbc.DBO_load_node_set_by_DB_id(id_set)
+        n_set = self.db_ctl.exec_op(op)
         self.assertEqual(len(n_set), len(id_set), 'incorrect result size')
 
     def test_partial_query_set_execution_success(self):