예제 #1
0
def setup_database_for_test():
    Database.init(debug=True)

    engine = Database.get_engine()
    connection = engine.connect()
    connection.execute("DROP SCHEMA IF EXISTS unittest CASCADE")
    connection.execute("CREATE SCHEMA unittest")
    connection.execute("SET search_path TO unittest, public")

    tx = connection.begin()
    try:
        # Create fresh schema
        with open(
                os.path.dirname(__file__) + "/../../src/rbb_server/schema.sql",
                'r') as sql_file:
            sql = sql_file.read()
        connection.execute(sql)

        # Insert all our testing data
        with open(os.path.dirname(__file__) + "/test-data.sql",
                  'r') as sql_file:
            sql = sql_file.read()

        statements = sql.split(";")
        for statement in statements:
            sql = statement.strip()
            if sql:
                connection.execute(sql)

        tx.commit()
    except Exception as e:
        tx.rollback()
        logging.error(e)
        print("ERROR: Cannot load unittest data into database!")
        raise e

    connection.close()

    # Make current session use the testing schema
    Database.get_session().execute("SET search_path TO unittest")
예제 #2
0
        def run(q, worker_name):
            rbb_server_test.database.init_database_connection_for_test()

            tasks = []
            number_of_collisions = 0
            while True:
                task = dequeue_task_inner(worker_name, "", "", admin_user)

                if isinstance(task, TaskDetailed):
                    tasks.append(task.description)
                    task.state = 100
                    put_task_inner(task.identifier, task, admin_user)
                else:
                    # Check if the queue is empty, if not then there was a collision
                    count = Database.get_engine().execute(
                        '''select count(uid) from task_queue where assigned_to='' '''
                    ).scalar()

                    if count > 0:
                        number_of_collisions += 1
                    else:
                        rbb_server_test.database.close_database()
                        q.put((tasks, number_of_collisions))
                        return