neo4j_test_util: gen_rand_data() -> DBO_random_data_generation

1 files changed, 42 insertions, 26 deletions

diff --git a/src/server-tests/neo4j_test_util.py b/src/server-tests/neo4j_test_util.py
index 90d9239b..6f49f17e 100644
--- a/src/server-tests/neo4j_test_util.py
+++ b/src/server-tests/neo4j_test_util.py
@@ -2,7 +2,48 @@ from random import choice
 import string
 import uuid
 
-from db_op import DBO_cypher_query
+from db_op import DB_op
+
+
+class DBO_random_data_generation(DB_op):
+
+    def __init__(self, lim_n=128, lim_r=256, prob_link_create=0.3):
+        """
+        generate random DB data
+        
+        @return: tuple consisting of the random node,link labels generated
+        """
+        assert 2 <= lim_n
+
+        super(DBO_random_data_generation, self).__init__()
+
+        self.n_label = rand_label()
+        self.r_label = rand_label()
+        q_arr = ['with 0 as _',  # TODO clean: foreach triggers SyntaxException: otherwise
+                 'foreach (rid in range(0,%d)' % (lim_n - 1),
+                 '|',
+                 'create (:%s {id: rid, n_attr_0:toInt(%d * rand())}))' % (self.n_label, lim_n)
+                ]
+
+        q = ' '.join(q_arr)
+        self.add_statement(q)
+
+        q_arr = ['match (s:%s),(d:%s)' % (self.n_label, self.n_label),
+                 'with s,d',
+                 'limit %d' % (lim_r - 1),
+                 'where rand() < %.2f' % (prob_link_create),
+                 'create (s)-[:%s {l_attr_0:toInt(%d * rand())}]->(d)' % (self.r_label, lim_r)]
+
+        q = ' '.join(q_arr)
+        self.add_statement(q)
+
+    @property
+    def node_set_label(self):
+        return self.n_label
+
+    @property
+    def link_set_label(self):
+        return self.r_label
 
 def rand_id():
     return str(uuid.uuid4())
@@ -21,31 +62,6 @@ def flush_db(db_ctl):
     db_ctl.exec_cypher_query('match (n) optional match (n)-[r]-() delete n,r')
 
 
-def gen_rand_data(db_ctl, lim_n=128, lim_r=256, prob_link_create=0.3):
-    """
-    generate random DB data
-    
-    @return: tuple consisting of the random node,link labels generated
-    """
-    assert 2 <= lim_n
-
-    n_label = rand_label()
-    r_label = rand_label()
-    q_arr = ['with 0 as _',  # TODO clean: foreach triggers SyntaxException: otherwise
-             'foreach (rid in range(0,%d)' % (lim_n - 1),
-             '|',
-             'create (:%s {id:rid, n_attr_0:toInt(%d * rand())}))' % (n_label, lim_n)
-            ]
-
-    q = ' '.join(q_arr)
-    op = DBO_cypher_query(q)
-    db_ctl.exec_op(op)
-
-    q_arr = ['match (s:%s),(d:%s)' % (n_label, n_label),
-             'with s,d',
-             'limit %d' % (lim_r - 1),
-             'where rand() < %.2f' % (prob_link_create),
-             'create (s)-[:%s {l_attr_0:toInt(%d * rand())}]->(d)' % (r_label, lim_r)]
 
     q = ' '.join(q_arr)
     op = DBO_cypher_query(q)