Python Registry.q方法代码示例

# 需要导入模块: from classes.Registry import Registry [as 别名]
# 或者: from classes.Registry.Registry import q [as 别名]

#.........这里部分代码省略.........
                    if field in ['referer', 'response_code', 'response_time', 'descr']:
                        row[field] = ''
                    elif field in ['spidered', 'size']:
                        row[field] = 0
                    elif field == 'who_add':
                        row[field] = 'human'
                    elif field == 'url':
                        raise WSException("URL row must have a 'url' key")

            for k in row.keys():
                if k not in [
                        'url', 'referer', 'response_code', 'response_time', 'who_add', 'spidered', 'size', 'descr'
                ]:
                    raise WSException("Key '{0}' must not be in url data".format(k))

            to_insert.append({
                'project_id': pid,
                "host_id": host_id,
                "hash": md5(row['url']),
                "url": row['url'],
                "referer": row['referer'],
                "response_code": row['response_code'],
                "response_time": row['response_time'],
                "when_add": int(time.time()),
                "who_add": row['who_add'],
                "spidered": row['spidered'],
                "size": row['size'],
                "descr": row['descr']

            })

            if len(to_insert)%50 == 0:
                self._db.insert_mass("urls", to_insert, 1)
                to_insert = []

        if len(to_insert):
            self._db.insert_mass("urls", to_insert, 1)

        return True

    def list_by_host_name(self, project_id, host, like=""):
        """ Get urls list by host name and project_id """
        host_id = HostsModel().get_id_by_name(project_id, host)

        like_expr = "" \
                    if not len(like.strip()) \
                    else " AND url LIKE '%{0}%' ".format(self._db.escape(like.strip()))

        return self._db.fetch_all(
            "SELECT url, response_code as code, response_time as time, when_add, who_add, descr, size FROM urls "
            "WHERE project_id = {0} AND host_id = {1} ".format(project_id, host_id) + like_expr +
            "ORDER BY url"
        )

    def list_by_host_name_for_spider(self, project_id, host):
        """ Get urls list by host name and project_id, but in special format for spider """
        host_id = HostsModel().get_id_by_name(project_id, host)

        return self._db.fetch_all(
            "SELECT url, response_code as code, response_time as time, when_add, who_add, descr FROM urls "
            "WHERE project_id = {0} AND host_id = {1} AND !spidered "
            "ORDER BY url".format(project_id, host_id)
        )

    def exists(self, project_id, host, url):
        """ Is url exists? """
        host_id = HostsModel().get_id_by_name(project_id, host)

        return self._db.fetch_one(
            "SELECT 1 FROM urls WHERE project_id = {0} AND host_id={1} AND hash = '{2}'"
            .format(project_id, host_id, md5(url))
        )

    def delete(self, project_id, host, url):
        """ Delete url from table """
        host_id = HostsModel().get_id_by_name(project_id, host)
        self._db.q(
            "DELETE FROM urls WHERE project_id = {0} AND host_id = {1} AND hash = {2} "
            .format(project_id, host_id, self._db.quote(md5(url)))
        )

    def update_url_field(self, project_id, host, url, field, value):
        """ Update custom field of current url """
        host_id = HostsModel().get_id_by_name(project_id, host)
        return self._db.update(
            "urls",
            {field: value},
            "hash = '{0}' AND project_id = {1} AND host_id = {2}".format(md5(url), project_id, host_id)
        )

    def update_url_field_mass(self, project_id, host, field, data):
        """ Mass update custom field of many urls """
        host_id = HostsModel().get_id_by_name(project_id, host)

        update = {}
        for row in data:
            case = "host_id = '{0}' AND `hash` = '{1}' ".format(host_id, md5(row['url']))
            update[case] = row['value']

        self._db.update_mass("urls", field, update)

# 需要导入模块: from classes.Registry import Registry [as 别名]
# 或者: from classes.Registry.Registry import q [as 别名]
class CommonTest(object):
    """ Common class tests """
    db = None

    def setup_class(self):
        open(wrpath + '/proxies.list', 'w').close()

        config = configparser.ConfigParser()
        config.read(os.getcwd() + '/' + 'config.ini')

        try:
            db = mysql.connector.connect(
                host=config['db']['host'],
                user=config['db']['user'],
                password=config['db']['pass'],
                database=config['db']['database']
            )
            db.autocommit = True
        except mysql.connector.errors.ProgrammingError as e:
            print " ERROR: Can`t connect to MySQL server! ({0})".format(str(e))
            exit(0)

        try:
            mc = MongoClient(host=config['mongo']['host'], port=int(config['mongo']['port']))
            mongo_collection = getattr(mc, config['mongo']['collection'])
        except pymongo.errors.ConnectionFailure as e:
            print " ERROR: Can`t connect to MongoDB server! ({0})".format(str(e))
            exit(0)

        R = Registry()
        R.set('config', config)
        R.set('db', db)
        R.set('mongo', mongo_collection)
        R.set('wr_path', wrpath)
        R.set('data_path', wrpath + '/data/')
        R.set('ndb',
              Database(config['db']['host'], config['db']['user'], config['db']['pass'], config['db']['database']))

        self.db = Registry().get('ndb')

    def setup(self):
        self.clear_db()

    def clear_db(self):
        """ Clearing db """
        self.db.q("TRUNCATE TABLE `hosts`")
        self.db.q("TRUNCATE TABLE `ips`")
        self.db.q("TRUNCATE TABLE `projects`")
        self.db.q("TRUNCATE TABLE `urls`")
        self.db.q("TRUNCATE TABLE `requests`")
        self.db.q("TRUNCATE TABLE `hosts_info`")
        self.db.q("TRUNCATE TABLE `urls_base`")
        self.db.q("TRUNCATE TABLE `urls_base_params`")
        self.db.q("TRUNCATE TABLE `cms`")
        self.db.q("TRUNCATE TABLE `cms_paths`")
        self.db.q("TRUNCATE TABLE `cms_paths_hashes`")

    def output_errors(self, out):
        """
        Check output for errors
        :param out: output
        :return:
        """
        assert out.find("Traceback") == -1
        assert "killed by time" not in out

    def _replace_config(self, name):
        """
        Change work config
        :param name: config name (*.ini)
        :return:
        """
        shutil.copyfile("{0}/config.ini".format(wrpath), "{0}/config.ini.bak".format(wrpath))
        shutil.copyfile("{0}/configs/{1}.ini".format(testpath, name), "{0}/config.ini".format(wrpath))

    def _restore_config(self):
        """ Restore old original config """
        shutil.move("{0}/config.ini.bak".format(wrpath), "{0}/config.ini".format(wrpath))

    def _run(self, config_name, run_params):
        """
        Run WS
        :param config_name: name of config for test
        :param run_params: params for run WS-process
        :return: output of WS
        """
        self._replace_config(config_name)
        os.chdir(wrpath)
        out = subprocess.check_output(run_params)
        self._restore_config()
        self.output_errors(out)

        return out

# 需要导入模块: from classes.Registry import Registry [as 别名]
# 或者: from classes.Registry.Registry import q [as 别名]
class CommonUnit(object):
    """ Common class for unit tests """
    db = None

    def setup_class(self):
        """ Prepare class for run tests """
        CURPATH = os.path.dirname(__file__) + "/"

        config = configparser.ConfigParser()
        config.read(CURPATH + 'config.ini')
        Registry().set('config', config)

        Registry().set('logger', LoggerMock())

        db = Database(
            config['main']['mysql_host'],
            config['main']['mysql_user'],
            config['main']['mysql_pass'],
            config['main']['mysql_dbname'],
        )
        Registry().set('db', db)

        self.db = Registry().get('db')  # type: Database

    def _clean_db(self):
        """ Clean tables for tests """
        self.db.q("TRUNCATE TABLE dicts")
        self.db.q("TRUNCATE TABLE dicts_groups")
        self.db.q("TRUNCATE TABLE hashes")
        self.db.q("TRUNCATE TABLE hashlists")
        self.db.q("TRUNCATE TABLE rules")
        self.db.q("TRUNCATE TABLE tasks")
        self.db.q("TRUNCATE TABLE tasks_groups")
        self.db.q("TRUNCATE TABLE task_works")
        self.db.q("TRUNCATE TABLE logs")

        self.db.update("algs", {'finder_insidepro_allowed': 0}, "id")

    def _add_hashlist(self, id=1, name='test', alg_id=3, have_salts=0,
                      status='ready', common_by_alg=0, parsed=1, last_finder_checked=0):
        """ Add hashlist record """
        self.db.insert(
            "hashlists",
            {
                'id': id,
                'name': name,
                'alg_id': alg_id,
                'have_salts': have_salts,
                'delimiter': '',
                'cracked': 0,
                'uncracked': 0,
                'errors': '',
                'parsed': parsed,
                'tmp_path': '',
                'status': status,
                'when_loaded': 0,
                'common_by_alg': common_by_alg,
                'last_finder_checked': last_finder_checked
            }
        )

    def _add_hash(self, hashlist_id=1, hash='', salt='', summ='', password='', cracked=0, id=None):
        """ Add hash record """
        self.db.insert(
            "hashes",
            {
                'id': id,
                'hashlist_id': hashlist_id,
                'hash': hash,
                'salt': salt,
                'password': password,
                'cracked': cracked,
                'summ': summ
            }
        )

    def _add_work_task(self, id=1, hashlist_id=1, task_id=1, status='wait', priority=0):
        """ Add work task record """
        self.db.insert(
            "task_works",
            {
                'id': id,
                'hashlist_id': hashlist_id,
                'task_id': task_id,
                'status': status,
                'priority': priority,
            }
        )

    def _add_task(self, id=1, name='task', group_id=1, type='dict', source=1):
        """ Add task record """
        self.db.insert(
            "tasks",
            {
                'id': id,
                'name': name,
                'group_id': group_id,
                'type': type,
                'source': source,
            }
#.........这里部分代码省略.........