Add parallel execution and shortcut notification to inspector design guideline
[doctor.git] / tests / common / utils.py
1 ##############################################################################
2 # Copyright (c) 2017 ZTE Corporation and others.
3 #
4 # All rights reserved. This program and the accompanying materials
5 # are made available under the terms of the Apache License, Version 2.0
6 # which accompanies this distribution, and is available at
7 # http://www.apache.org/licenses/LICENSE-2.0
8 ##############################################################################
9 import json
10 import os
11 import paramiko
12 import re
13
14
15 def load_json_file(full_path):
16     """Loads JSON from file
17     :param target_filename:
18     :return:
19     """
20     if not os.path.isfile(full_path):
21         raise Exception('File(%s) does not exist' % full_path)
22
23     with open(full_path, 'r') as file:
24         return json.load(file)
25
26
27 def write_json_file(full_path, data):
28     """write JSON from file
29     :param target_filename:
30     :return:
31     """
32
33     with open(full_path, 'w+') as file:
34         file.write(json.dumps(data))
35
36
37 def match_rep_in_file(regex, full_path):
38     if not os.path.isfile(full_path):
39         raise Exception('File(%s) does not exist' % full_path)
40
41     with open(full_path, 'r') as file:
42         for line in file:
43             result = re.search(regex, line)
44             if result:
45                 return result
46
47     return None
48
49
50 class SSHClient(object):
51     def __init__(self, ip, username, password=None, pkey=None,
52                  key_filename=None, log=None, look_for_keys=False,
53                  allow_agent=False):
54         self.client = paramiko.SSHClient()
55         self.client.load_system_host_keys()
56         self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
57         self.client.connect(ip, username=username, password=password,
58                             pkey=pkey, key_filename=key_filename,
59                             look_for_keys=look_for_keys,
60                             allow_agent=allow_agent)
61         self.log = log
62
63     def __del__(self):
64         self.client.close()
65
66     def ssh(self, command):
67         if self.log:
68             self.log.info("Executing: %s" % command)
69         stdin, stdout, stderr = self.client.exec_command(command)
70         ret = stdout.channel.recv_exit_status()
71         output = list()
72         for line in stdout.read().splitlines():
73             output.append(line.decode('utf-8'))
74         if ret:
75             if self.log:
76                 self.log.info("*** FAILED to run command %s (%s)" % (command, ret))
77             raise Exception(
78                 "Unable to run \ncommand: %s\nret: %s"
79                 % (command, ret))
80         if self.log:
81             self.log.info("*** SUCCESSFULLY run command %s" % command)
82         return ret, output
83
84     def scp(self, source, dest, method='put'):
85         if self.log:
86             self.log.info("Copy %s -> %s" % (source, dest))
87         ftp = self.client.open_sftp()
88         if method == 'put':
89             ftp.put(source, dest)
90         elif method == 'get':
91             ftp.get(source, dest)
92         ftp.close()
93
94
95 def run_async(func):
96     from threading import Thread
97     from functools import wraps
98
99     @wraps(func)
100     def async_func(*args, **kwargs):
101         thread = Thread(target=func, args=args, kwargs=kwargs)
102         thread.start()
103         return thread
104
105     return async_func