1from librabbitmq import Connection
2import argparse
3import paramiko
4import sys
5import time
6import os
7import shutil
8
9python_exe = "python"
10if os.system("grep \'centos\' /etc/issue -i -q") == 0:
11    python_exe = "python2.7"
12
13
14def get_ssh_client(ip, username=None, password=None, timeout=10):
15    client = None
16    try:
17        ip = ip.split(':')[0]
18        client = paramiko.SSHClient()
19        client.load_system_host_keys()
20        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
21        username = username or cfg.SSH_USER
22        password = password or cfg.SSH_PASSWORD
23        client.connect(ip, username=username, password=password, timeout=timeout)
24        print("Successfully SSHed to {0}".format(ip))
25    except Exception as ex:
26        print ex
27        sys.exit(1)
28    return client
29
30
31def get_sftp_client(ip, username=None, password=None,):
32    try:
33        ip = ip.split(':')[0]
34        trans = paramiko.Transport((ip))
35        username = username or cfg.SSH_USER
36        password = password or cfg.SSH_PASSWORD
37        trans.connect(username=username, password=password)
38        print ("SFTPing to {0}".format(ip))
39        sftp_client = paramiko.SFTPClient.from_transport(trans)
40        return sftp_client
41    except Exception as ex:
42        print ex
43        sys.exit(1)
44
45
46def kill_process(ssh_client, process_name):
47    print "Killing {0}".format(process_name)
48    _, stdout, _ = ssh_client.exec_command("pgrep -f {0}".format(process_name))
49    for pid in stdout.readlines():
50        ssh_client.exec_command("kill -9 {0}".format(pid.split()[0]))
51
52
53def start_process(ssh_client, process_name, cmd):
54    print ("Starting {0}...".format(process_name))
55    ssh_client.exec_command(cmd)
56    time.sleep(5)
57    _, stdout, _ = ssh_client.exec_command("pgrep {0}".format(process_name))
58    print ("{0} is running with pid {1}".format(process_name, stdout.readlines()[0]))
59
60
61def start_rabbitmq():
62    vhost_present = False
63    tries = 1
64    print("\n##### Setting up RabbitMQ @ {0} #####".format(cfg.RABBITMQ_IP))
65    rabbitmq_client = get_ssh_client(cfg.RABBITMQ_IP)
66
67    _, stdout, _ = rabbitmq_client.exec_command("ps aux|grep rabbitmq|grep -v grep|awk \'{print $2}\'")
68    print ("Killing existing RabbitMQ process ...")
69    for pid in stdout:
70        if pid == "":
71            continue
72        rabbitmq_client.exec_command("sudo kill -9 {0}".format(pid))
73    if cfg.RABBITMQ_LOG_LOCATION is not "":
74        print("Deleting RabbitMQ logs from {0}".format(cfg.RABBITMQ_LOG_LOCATION))
75        rabbitmq_client.exec_command("rm -rf {0}/*.*".format(cfg.RABBITMQ_LOG_LOCATION))
76    print ("Starting RabbitMQ ...")
77    rabbitmq_client.exec_command("screen -dmS rabbitmq sh -c \'sudo rabbitmq-server start; exec bash;\'")
78    time.sleep(20)
79    _, stdout, _ = rabbitmq_client.exec_command("sudo rabbitmqctl status")
80    for line in stdout.readlines():
81        sys.stdout.write(line)
82    print("Rabbitmq has been restarted and is now running!")
83    _, stdout, _ = rabbitmq_client.exec_command("sudo rabbitmqctl list_vhosts")
84    for line in stdout.readlines():
85        if not vhost_present:
86            if cfg.CB_CLUSTER_TAG in line:
87                vhost_present = True
88        sys.stdout.write(line)
89    if not vhost_present :
90        print ("Adding vhost {0} and setting permissions".format(cfg.CB_CLUSTER_TAG))
91        rabbitmq_client.exec_command("sudo rabbitmqctl add_vhost {0}".format(cfg.CB_CLUSTER_TAG))
92        rabbitmq_client.exec_command("sudo rabbitmqctl set_permissions -p {0} guest '.*' '.*' '.*'".format(cfg.CB_CLUSTER_TAG))
93        _, stdout, _ = rabbitmq_client.exec_command("sudo rabbitmqctl list_vhosts")
94        for line in stdout.readlines():
95            sys.stdout.write(line)
96    time.sleep(30)
97    while True:
98        try:
99            tries += 1
100            Connection(host=cfg.RABBITMQ_IP, userid="guest", password="guest", virtual_host=cfg.CB_CLUSTER_TAG)
101            print("Connected to RabbitMQ vhost")
102            break
103        except Exception as e:
104            print e
105            if tries <= 5:
106                print("Retrying connection {0}/5 ...".format(tries))
107                rabbitmq_client.exec_command("sudo rabbitmqctl delete_vhost {0}".format(cfg.CB_CLUSTER_TAG))
108                rabbitmq_client.exec_command("sudo rabbitmqctl add_vhost {0}".format(cfg.CB_CLUSTER_TAG))
109                rabbitmq_client.exec_command("sudo rabbitmqctl set_permissions -p {0} guest '.*' '.*' '.*'".format(cfg.CB_CLUSTER_TAG))
110                time.sleep(30)
111                continue
112            sys.exit(1)
113    rabbitmq_client.close()
114
115
116def start_worker(worker_ip):
117    print("##### Setting up Celery Worker @ {0} #####".format(worker_ip))
118    worker_client = get_ssh_client(worker_ip)
119
120#    # Update Worker's testrunner repository
121#    repo_dir = os.path.dirname(os.path.dirname(cfg.WORKER_PYSYSTESTS_PATH))
122#    worker_client.exec_command("rm -rf {0}; mkdir -p {0}".format(repo_dir))
123#    worker_client.exec_command("cd {0}; git clone https://github.com/couchbase/testrunner.git".format(repo_dir))
124
125    # Copy testcfg.py file to all workers
126    worker_client.open_sftp().put("./testcfg.py", os.path.join(cfg.WORKER_PYSYSTESTS_PATH, "testcfg.py"))
127
128    # kill celery,remove screenlog
129    kill_process(worker_client, "celery")
130    worker_client.exec_command("screen -ls | grep \'celery\' | awk '{print $1}' | xargs -i screen -X -S {} quit")
131    worker_client.exec_command("screen -wipe")
132    worker_client.exec_command("rm -rf {0}/screenlog.0".format(cfg.WORKER_PYSYSTESTS_PATH))
133
134    # memcached
135    kill_process(worker_client, "memcached")
136    cmd = "memcached -u couchbase -d -l {0} -p 11911".format(worker_ip)
137    start_process(worker_client, "memcached", cmd)
138
139    print("Starting celery worker...")
140
141    _, out, _ = worker_client.exec_command("celery --version")
142    celery_param = ""
143    for line in out:
144        if "3.1.16" in line:
145            celery_param = "-Ofair"
146            print "Celery version: {0} is installed, running it with {1} param".format(line, celery_param)
147            break
148
149    if worker_ip == cfg.WORKERS[0]:
150        _, stdout, _ = worker_client.exec_command("cd {0}; pwd; export C_FORCE_ROOT=1;screen -dmS celery -L sh -c  \ "
151        "\'celery worker -c 8 -A app -B -l ERROR {1} --purge -I app.init; exec bash;\'".format(cfg.WORKER_PYSYSTESTS_PATH, celery_param))
152    else:
153        _, stdout, _ = worker_client.exec_command("cd {0}; pwd; screen -dmS celery -L sh -c \
154         \'celery worker -c 16 -A app -l ERROR {1} -I app.init; exec bash;\'".format(cfg.WORKER_PYSYSTESTS_PATH, celery_param))
155    time.sleep(20)
156    #read_screenlog(worker_ip, cfg.WORKER_PYSYSTESTS_PATH, stop_if_EOF=True)
157    worker_client.close()
158
159
160def start_seriesly():
161    print("##### Setting up Seriesly @ {0} #####".format(cfg.SERIESLY_IP))
162    cbmonitor_client = get_ssh_client(cfg.SERIESLY_IP)
163
164    kill_process(cbmonitor_client, "seriesly")
165    if cfg.SERIESLY_DB_LOCATION is not "":
166        print("Deleting old Seriesly db files from {0}".format(cfg.SERIESLY_DB_LOCATION))
167        cbmonitor_client.exec_command("rm -rf {0}/*.*".format(cfg.SERIESLY_DB_LOCATION))
168    # kill all existing screens
169    cbmonitor_client.exec_command("screen -ls | grep \'seriesly\' | awk \'{print $1}\' | xargs -i screen -X -S {} quit")
170    cbmonitor_client.exec_command("screen -ls | grep \'webapp\' | awk \'{print $1}\' | xargs -i screen -X -S {} quit")
171    cbmonitor_client.exec_command("screen -ls | grep \'ns_collector\' | awk \'{print $1}\' | xargs -i screen -X -S {} quit")
172    cbmonitor_client.exec_command("screen -ls | grep \'atop_collector\' | awk \'{print $1}\' | xargs -i screen -X -S {} quit")
173    cbmonitor_client.exec_command("rm -rf {0}/screenlog.0".format(cfg.CBMONITOR_HOME_DIR))
174    # screen 1 - start seriesly
175    start_cmd = "screen -dmS seriesly -L sh -c \'cd {0}; ./seriesly; exec bash;\'".format(cfg.SERIESLY_LOCATION)
176    start_process(cbmonitor_client, "seriesly", start_cmd)
177
178
179def fix_sample_cfg(ssh_client):
180    # fix sample.cfg file
181    cfg_file_path = os.path.join(cfg.CBMONITOR_HOME_DIR, "sample.cfg")
182    ssh_client.exec_command("sed -i 's/.*host_port.*/host_port = {0}:8000/' {1}".format(cfg.SERIESLY_IP, cfg_file_path))
183    ssh_client.exec_command("sed -i 's/.*host .*/host = {0}/' {1}".format(cfg.SERIESLY_IP, cfg_file_path))
184    ssh_client.exec_command("sed -i 's/.*master_node.*/master_node = {0}/' {1}".format(cfg.COUCHBASE_IP, cfg_file_path))
185
186    if cfg.COUCHBASE_OS == "windows":
187        ssh_client.exec_command("sed -i 's/.*ssh_username.*/ssh_username = {0}/' {1}".format(cfg.COUCHBASE_SSH_USER, cfg_file_path))
188        ssh_client.exec_command("sed -i 's/.*ssh_password.*/ssh_password = {0}/' {1}".format(cfg.COUCHBASE_SSH_PASSWORD, cfg_file_path))
189
190
191def start_cbmonitor():
192    print("\n##### Setting up CBMonitor @ {0} #####".format(cfg.SERIESLY_IP))
193    cbmonitor_client = get_ssh_client(cfg.SERIESLY_IP)
194    # screen 2 - start webserver
195    kill_process(cbmonitor_client, "webapp")
196    start_cmd = "cd {0}; screen -dmS webapp -L sh -c \'./bin/webapp add-user -S;./bin/webapp syncdb; \
197     ./bin/webapp runserver {1}:8000; exec bash;\'".format(cfg.CBMONITOR_HOME_DIR, cfg.SERIESLY_IP)
198    start_process(cbmonitor_client, "webapp", start_cmd)
199
200    # screen 3 - start ns_collector
201    fix_sample_cfg(cbmonitor_client)
202    kill_process(cbmonitor_client, "ns_collector")
203    start_cmd = "cd {0}; screen -dmS ns_collector -L sh -c \'./bin/ns_collector sample.cfg; exec bash;\'".format(cfg.CBMONITOR_HOME_DIR)
204    start_process(cbmonitor_client, "ns_collector", start_cmd)
205
206    # screen 4 - start atop_collector
207    kill_process(cbmonitor_client, "atop_collector")
208    start_cmd = "cd {0}; screen -dmS atop_collector -L sh -c \'./bin/atop_collector sample.cfg; exec bash;\'".format(cfg.CBMONITOR_HOME_DIR)
209    start_process(cbmonitor_client, "atop_collector", start_cmd)
210    read_screenlog(cfg.SERIESLY_IP, cfg.CBMONITOR_HOME_DIR, stop_if_EOF=True, lines_to_read=100)
211    cbmonitor_client.close()
212
213
214def read_screenlog(ip, screenlog_dir, retry=10, stop_if_EOF=False, lines_to_read=20000):
215    line = ""
216    line_count = 0
217    last_pos = 0
218    transport_client = get_sftp_client(ip)
219    screen_log = "{0}/screenlog.0".format(screenlog_dir)
220    op_file = transport_client.open(screen_log, 'r')
221    while "Test Complete" not in line and line_count < lines_to_read:
222        op_file.seek(last_pos)
223        line = op_file.readline()
224        last_pos = op_file.tell()
225        if line is not None and line is not "":
226            sys.stdout.write(line)
227            line_count += 1
228        else:
229            #Reached EOF, will retry after 'retry' secs
230            if stop_if_EOF:
231                break
232            time.sleep(retry)
233    op_file.close()
234    transport_client.close()
235
236
237def run_setup():
238    # kick off the setup test
239    print("\n##### Starting cluster setup from {0} #####".format(cfg.SETUP_JSON))
240    worker_client = get_ssh_client(cfg.WORKERS[0])
241    # Import templates if needed
242    for template in cfg.SETUP_TEMPLATES:
243        print ("Importing document template {0}...".format(template.split('--')[1].split('--')[0]))
244        temp = "{0} cbsystest.py import template {1}".format(python_exe, template)
245        print temp
246        _, stdout, _ = worker_client.exec_command("cd {0}; {1} cbsystest.py import template {2} --cluster {3}".
247                                                  format(cfg.WORKER_PYSYSTESTS_PATH, python_exe, template, cfg.CB_CLUSTER_TAG))
248        for line in stdout.readlines():
249            print line
250    print ("Running test ...")
251    _, stdout, _ = worker_client.exec_command("cd {0}; {1} cbsystest.py run test --cluster \'{2}\' --fromfile \'{3}\'".
252                                              format(cfg.WORKER_PYSYSTESTS_PATH, python_exe, cfg.CB_CLUSTER_TAG, cfg.SETUP_JSON))
253    read_screenlog(cfg.WORKERS[0], cfg.WORKER_PYSYSTESTS_PATH)
254    worker_client.close()
255
256def run_test():
257    print "\n##### Starting system test #####"
258    start_worker(cfg.WORKERS[0])
259    # import doc template in worker
260    worker_client = get_ssh_client(cfg.WORKERS[0])
261    for template in cfg.TEST_TEMPLATES:
262        print ("Importing document template {0}...".format(template.split('--')[1].split('--')[0]))
263        temp = "{0} cbsystest.py import template {1}".format(python_exe, template)
264        print temp
265        _, stdout, _ = worker_client.exec_command("cd {0}; {1} cbsystest.py import template {2} --cluster {3}".
266                                                  format(cfg.WORKER_PYSYSTESTS_PATH, python_exe, template, cfg.CB_CLUSTER_TAG))
267        for line in stdout.readlines():
268            print line
269    # Start sys test
270    print ("Starting system test from {0}...".format(cfg.TEST_JSON))
271    _, stdout, _ = worker_client.exec_command("cd {0}; {1} cbsystest.py run test --cluster \'{2}\' --fromfile \'{3}\'".
272                                              format(cfg.WORKER_PYSYSTESTS_PATH, python_exe, cfg.CB_CLUSTER_TAG, cfg.TEST_JSON))
273    time.sleep(5)
274    for line in stdout.readlines():
275        sys.stdout.write(line)
276    read_screenlog(cfg.WORKERS[0], cfg.WORKER_PYSYSTESTS_PATH)
277    worker_client.close()
278
279def pre_install_check():
280    try:
281        print("##### Pre-install inspection #####")
282        print("Inspecting Couchbase server VMs ...")
283        for vm_ip in cfg.CLUSTER_IPS:
284            if cfg.COUCHBASE_OS == "windows":
285                vm_client = get_ssh_client(vm_ip, cfg.COUCHBASE_SSH_USER, cfg.COUCHBASE_SSH_PASSWORD)
286            else:
287                vm_client = get_ssh_client(vm_ip)
288            vm_client.close()
289        print ("Inspecting RabbitMQ ...")
290        rabbitmq = get_ssh_client(cfg.RABBITMQ_IP)
291        rabbitmq.close()
292        print ("Inspecting Worker ...")
293        worker = get_ssh_client(cfg.WORKERS[0])
294        worker.close()
295        print ("Inspecting CBMonitor ...")
296        cbmonitor = get_ssh_client(cfg.SERIESLY_IP)
297        cbmonitor.close()
298        print("Inspection complete!")
299    except Exception as e:
300        print e
301        sys.exit()
302
303def upload_stats():
304    print "\n##### Uploading stats to CBFS #####"
305    worker_client = get_ssh_client(cfg.WORKERS[0])
306    push_stats_cmd = "cd {0}; {1} tools/push_stats.py  --version {2} --build {3} --spec {4} \
307    --name {5} --cluster {6}".format(cfg.WORKER_PYSYSTESTS_PATH, python_exe, args['build'].split('-')[0], args['build'].split('-')[1],
308    cfg.TEST_JSON, cfg.TEST_JSON[cfg.TEST_JSON.rfind('/') + 1 : cfg.TEST_JSON.find('.')] , cfg.CB_CLUSTER_TAG)
309    print ("Executing {0}".format(push_stats_cmd))
310    _, stdout, _ = worker_client.exec_command(push_stats_cmd)
311    time.sleep(30)
312    for line in stdout.readlines():
313        print line
314    worker_client.close()
315
316def install_couchbase():
317    print("Installing version {0} Couchbase on servers ...".format(args['build']))
318    install_cmd = "cd ..; {0} scripts/install.py -i {1} -p product=cb,version={2},parallel=true,{3}".\
319                    format(python_exe, cfg.CLUSTER_INI, args['build'], args['params'])
320    print("Executing : {0}".format(install_cmd))
321    os.system(install_cmd)
322    if cfg.CLUSTER_RAM_QUOTA != "":
323        os.system("curl -d memoryQuota={0} \"http://{1}:{2}@{3}:8091/pools/default\"".
324                  format(cfg.CLUSTER_RAM_QUOTA, cfg.COUCHBASE_USER, cfg.COUCHBASE_PWD, cfg.CLUSTER_IPS[0]))
325    for ip in cfg.CLUSTER_IPS:
326        os.system("curl -X POST -d \'ale:set_loglevel(xdcr_trace, debug).\' \"http://{0}:{1}@{2}:8091/diag/eval\"".
327                  format(cfg.COUCHBASE_USER, cfg.COUCHBASE_PWD, ip))
328    time.sleep(60)
329
330def warn_skip(task):
331    print("\nWARNING : Skipping {0}\n".format(task))
332    return True
333
334def run(args):
335    exlargs = args['exclude']
336
337    # Pre-install check
338    ("inspect" in exlargs) and warn_skip("Inspection") or pre_install_check()
339
340    # Install Couchbase
341    ("install" in exlargs) and warn_skip("Installation") or install_couchbase()
342
343    # Setup RabbitMQ
344    ("rabbitmq" in exlargs) and warn_skip("RabbitMQ") or start_rabbitmq()
345
346    # Setup Seriesly
347    ("seriesly" in exlargs) and warn_skip("Seriesly") or start_seriesly()
348
349    # Start workers
350    ("worker" in exlargs) and warn_skip("Celery Worker setup") or\
351                               [start_worker(ip) for ip in cfg.WORKERS]
352
353    # Cluster-setup/create buckets, set RAM quota
354    ("setup" in exlargs) and warn_skip("Cluster setup") or run_setup()
355
356    # Start cbmonitor
357    ("cbmonitor" in exlargs) and warn_skip("CBMonitor") or start_cbmonitor()
358
359    # Run test
360    ("systest" in exlargs) and warn_skip("System Test") or run_test()
361
362    # Upload stats
363    ("stats" in exlargs) and warn_skip("Uploading Stats to CBFS") or upload_stats()
364
365    print("\n############################# Execution Complete! #################################")
366
367
368if __name__ == "__main__":
369    parser = argparse.ArgumentParser(description="Tool for running system tests \
370                 \nUsage: python runsystest.py --build 3.0.0-355 \
371                                              --testcfg xdcr/testcfg_source.py \
372                                              --params upr=true,xdcr_upr=false \
373                                              --exclude install,seriesly,worker,cbmonitor,cluster,systest,stats")
374    parser.add_argument("--build", help="required param: build-version for system test to run on", required=True)
375    parser.add_argument("--testcfg", default="testcfg.py", help="required param: location of testcfg file in testcfg dir ")
376    parser.add_argument("--params", help="optional param: additional build params eg:vbuckets=1024,upr=true,xdcr_upr=false",
377                        required=False)
378    parser.add_argument("--exclude",
379                            nargs='+',
380                            default="",
381                            help="optional param: inspect install rabbitmq seriesly worker cbmonitor setup systest stats",
382                            required=False)
383
384    try:
385        args = vars(parser.parse_args())
386        testcfg = args['testcfg']
387        if os.path.basename(os.path.abspath(os.getcwd())) != 'pysystests':
388            raise Exception("Run script from testrunner/pysystests folder, current folder is: %s" % os.getcwd())
389        shutil.copy(testcfg, "./testcfg.py")
390        print "Copied {0} to {1}/testcfg.py".format(testcfg, os.getcwd())
391        cfg = __import__("testcfg")
392        run(args)
393    except Exception as e:
394        print e
395        raise
396