import os
from glob import glob
import shutil
import pprint

import traceback
import warnings
import json
import yaml

import intake

import dask
import xarray as xr

from toolz import curry

from . config import cache_catalog_dir, cache_catalog_prefix
from . registry import _DERIVED_VAR_REGISTRY, _QUERY_DEPENDENT_OP_REGISTRY

os.makedirs(cache_catalog_dir, exist_ok=True)

class Collection(object):
    """
    Catalog + Query + Operation = Unique Dataset

    Extend an intake-esm catalog with the ability to:
      (1) Compute derived variables combining multipe catalog entries;
      (2) Apply postprocessing operations to returned datasets;
      (3) Optionally cache the results and curate a discrete catalog of processed assets;
          skip computation on subsequent calls for existing cataloged assets.

    Parameters
    ----------

    name : str
      Name of this collection.

    catalog : str
      JSON file defining `intake-esm` catalog

    query : dict
      Dictionary of keyword, value pairs defining a query subsetting the `catalog`.

    postproccess : callable, list of callables, optional
      Call these functions on each dataset following aggregation.

    postproccess_kwargs: dict, list of dict's, optional
       Keyword arguments to `postprocess` functions.
       **Note** These arguments must be serializable to yaml.

    persist : bool, optional
       Persist datasets to disk.

    cache_dir : str, optional
       Directory for storing cache files.

    cache_format : str
       Format in which to store cache files; can be "nc" or "zarr"

    kwargs : dict
       Defaults for keyword arguments to pass to `intake_esm.core.esm_datastore.to_dataset_dict`.
       (Note: these can be updated/overridden for each call to self.dsets below.)
    """
    def __init__(self,
                 name,
                 esm_collection_json,
                 query,
                 postproccess=[],
                 postproccess_kwargs=[],
                 persist=False,
                 cache_dir='.',
                 cache_format='nc',
                 **kwargs,
                ):

        self.name = name

        # TODO: what should we do if "variable" is in the query?
        #       for now, just remove it
        variable = query.get('variable', None)
        if variable is not None:
            warnings.warn('removed "variable" key from query')

        # TODO: accept multiple catalogs and concatenate them into one?
        self.catalog = intake.open_esm_datastore(esm_collection_json).search(**query)

        # setup cache info
        assert cache_format in ['nc', 'zarr'], f'unsupported format {cache_format}'
        self._format = cache_format
        self.persist = persist
        self.cache_dir = cache_dir
        if self.persist and not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
        self._open_cache_kwargs = dict() # TODO: get this from config?
        if self._format == 'zarr' and 'consolidated' not in self._open_cache_kwargs:
            self._open_cache_kwargs['consolidated'] = True

        self._to_dsets_kwargs_defaults = kwargs

        # setup operators attrs
        if not isinstance(postproccess, list):
            postproccess = [postproccess]
        self.operators = postproccess

        if not postproccess_kwargs:
            self.ops_kwargs = [{} for op in self.operators]
        else:
            if not isinstance(postproccess_kwargs, list):
                postproccess_kwargs = [postproccess_kwargs]
            assert len(postproccess_kwargs) == len(postproccess), 'mismatched ops/ops_kwargs'
            self.ops_kwargs = postproccess_kwargs

        # pull out "prepocess"
        preprocess_name = None if 'preprocess' not in kwargs else kwargs['preprocess'].__name__

        # build origins dict
        # TODO: save more about the function, including the source code
        #       rather than simply the name
        self.origins_dict = dict(
            name=name,
            esm_collection=esm_collection_json,
            query=query,
            preprocess=preprocess_name,
            operators=[op.__name__ for op in self.operators],
            operator_kwargs=self.ops_kwargs,
        )

        # assemble database of existing cache's
        self._assemble_cache_db()

    def _assemble_cache_db(self):
        """loop over yaml files in cache dir; find ones that match"""
        self._cache_files = {}
        if not self.persist:
            return

        cache_id_files = self._find_cache_id_files()

        if cache_id_files:
            for file in cache_id_files:
                try:
                    with open(file, 'r') as fid:
                        cache_id = yaml.load(fid, Loader=yaml.Loader)
                except: # TODO: be more informative here
                    print(f'cannot read cache: {file}...skipping')
                    print('skipping')

                cache_origins_dict = {
                    k: cache_id[k]
                    for k in self.origins_dict.keys()
                    if k in cache_id
                }
                if cache_origins_dict == self.origins_dict:
                    variable = cache_id['variable']
                    key = cache_id['key']
                    if variable not in self._cache_files:
                        self._cache_files[variable] = {}
                    self._cache_files[variable][key] = cache_id['asset']

    def to_dataset_dict(self, variable, compute=True, clobber=False,
                        prefer_derived=False, **kwargs):
        """
        Get dataset_dicts for assets.

        Parameters
        ----------

        variable : str or list
          The variable or list of variables to read and process.

        compute : bool, optional (default=True)
          Specify whether to call compute on the final datasets.

        clobber : bool, optional (default=False)
          Remove and recreate any existing cache.

        kwargs : dict, optional
          Keyword arguments to pass to `intake_esm.core.esm_datastore.to_dataset_dict`

        """

        # TODO: `variable` could be more flexible,
        # i.e. it could be a query dict with a require "variable" key

        if isinstance(variable, list):
            dsets_list = [
                self.to_dataset_dict(v, compute, clobber, prefer_derived, **kwargs)
                for v in variable
            ]
            keys = list(set([k for dsets in dsets_list for k in dsets.keys()]))
            dsets_merged = {}
            for key in keys:
                ds_list = []
                for dsets in dsets_list:
                    if key in dsets:
                        ds_list.append(dsets[key])
                dsets_merged[key] = xr.merge(ds_list)
            return dsets_merged

        # TODO: check for lock file, wait if present
        if self._cache_exists(variable, clobber):
            return self._read_cache(variable)

        else:
            # TODO: set lock file
            to_dsets_kwargs = kwargs.copy()
            to_dsets_kwargs.update(self._to_dsets_kwargs_defaults)
            # TODO: optionally spin up a cluster here
            return self._generate_dsets(
                variable, compute, prefer_derived, **to_dsets_kwargs
            )
            # TODO: release lock

    def _generate_dsets(self, variable, compute, prefer_derived, **kwargs):
        """Do the computation necessary to make `dsets`"""

        # check for variable in catalog
        # TODO: we're handling empty results below:
        #       suppress intake-esm warning or change logic
        catalog_subset_var = self.catalog.search(variable=variable)

        # check for variable in derived registry
        is_derived = variable in _DERIVED_VAR_REGISTRY

        if len(catalog_subset_var):
            if is_derived and not prefer_derived:
                warnings.warn(
                    f'found variable "{variable}" in catalog and derived_var registry'
                )
                is_derived = False
        else:
            if not is_derived:
                raise ValueError(f'variable not found {variable}')

        if is_derived:
            derived_var_obj = _DERIVED_VAR_REGISTRY[variable]
            query_vars = derived_var_obj.dependent_vars
            catalog_subset_var = self.catalog.search(variable=query_vars)

        dsets = catalog_subset_var.to_dataset_dict(**kwargs)
        key_info = _intake_esm_get_keys_info(catalog_subset_var)

        if is_derived:
            for key, ds in dsets.items():
                dsets[key] = derived_var_obj(ds)

        for key in dsets.keys():
            for op, kw in zip(self.operators, self.ops_kwargs):
                if hash(op) in _QUERY_DEPENDENT_OP_REGISTRY:
                    op_obj = _QUERY_DEPENDENT_OP_REGISTRY[hash(op)]
                    dsets[key] = op_obj(
                        dsets[key],
                        query_dict=key_info[key],
                        **kw,
                    )
                else:
                    dsets[key] = op(dsets[key], **kw)

        if compute:
            dsets = {k: ds.compute() for k, ds in dsets.items()}

        if self.persist:
            self._make_cache(dsets, variable)

        return dsets

    def _make_cache(self, dsets, variable):
        """write cache file"""

        cache_files = {variable: {}}
        for key, ds in dsets.items():

            cache_id_dict = self.origins_dict.copy()
            cache_id_dict['key'] = key
            cache_id_dict['variable'] = variable
            cache_id_dict['asset'] = self._gen_cache_file_name(key, variable)

            cache_id_file = self._gen_cache_id_file_name(key, variable)

            if self._format == 'nc':
                ds.to_netcdf(cache_id_dict['asset'])
            elif self._format == 'zarr':
                ds.to_zarr(cache_id_dict['asset'], mode='w', consolidated=True)

            with open(cache_id_file, 'w') as fid:
                yaml.dump(cache_id_dict, fid)

            cache_files[variable][key] = cache_id_dict['asset']

        self._cache_files.update(cache_files)

    def _read_cache(self, variable):
        """read cache files"""
        dsets = {}
        for key, asset in self._cache_files[variable].items():
            if self._format == 'nc':
                with xr.open_dataset(asset, **self._open_cache_kwargs) as ds:
                    dsets[key] = ds
            elif self._format == 'zarr':
                with xr.open_zarr(asset, **self._open_cache_kwargs) as ds:
                    dsets[key] = ds
        return dsets

    def _cache_exists(self, variable, clobber):
        """determine if cache files exist (or clobber them)"""

        if variable in self._cache_files:
            if clobber:
                for asset in self._cache_files[variable].values():
                    if os.path.exists(asset):
                        self._remove_asset(asset)
                return False
            else:
                return all(
                    [os.path.exists(asset) for asset in self._cache_files[variable].values()]
                )
        else:
            return False

    def _remove_asset(self, asset):
        """delete asset from disk"""
        if not os.path.exists(asset):
            return
        if self._format == 'nc':
            os.remove(asset)
        elif self._format == 'zarr':
            shutil.rmtree(asset)

    def _gen_cache_file_name(self, key, variable):
        """generate a file cache name"""
        # TODO: accept a user-provided callable to generate human-readable
        #       file name
        token = dask.base.tokenize(self.origins_dict, key, variable)
        return f'{self.cache_dir}/{token}.{self._format}'

    def _gen_cache_id_file_name(self, key, variable):
        """generate a unique cache file name"""
        token = dask.base.tokenize(self.origins_dict, key, variable)
        return f'{cache_catalog_dir}/{cache_catalog_prefix}-{token}.yml'

    def _find_cache_id_files(self):
        return sorted(glob(f'{cache_catalog_dir}/{cache_catalog_prefix}-*.yml'))

    def __repr__(self):
        return pprint.pformat(self.origins_dict, indent=2, width=1)

class derived_var(object):
    """
    Support computation of variables that depend on multiple variables,
    i.e., "derived vars"
    """
    def __init__(self, dependent_vars, func):
        self.dependent_vars = dependent_vars
        self._callable = func

    def __call__(self, ds, **kwargs):
        """call the function to compute derived var"""
        self._ensure_variables(ds)
        return self._callable(ds, **kwargs)

    def _ensure_variables(self, ds):
        """ensure that required variables are present"""
        missing_var = set(self.dependent_vars) - set(ds.variables)
        if missing_var:
            raise ValueError(f'Variables missing: {missing_var}')

class query_dependent_op(object):
    """
    Support calling functions that depend on the values of the query.
    """
    def __init__(self, func, query_keys):
        self._callable = func
        self._query_keys = query_keys

    def __call__(self, ds, query_dict, **kwargs):
        """call function with query keys added to keyword args"""
        kwargs.update({k: query_dict[k] for k in self._query_keys})
        return self._callable(ds, **kwargs)

@curry
def register_derived_var(func, varname, dependent_vars):
    """register a function for computing derived variables"""
    if varname in _DERIVED_VAR_REGISTRY:
        warnings.warn(
            f'overwriting derived variable "{varname}" definition'
        )

    _DERIVED_VAR_REGISTRY[varname] = derived_var(
        dependent_vars, func,
    )
    return func

@curry
def register_query_dependent_op(func, query_keys):
    """register a function for computing derived variables"""
    func_hash = hash(func)
    if func_hash in _QUERY_DEPENDENT_OP_REGISTRY:
        warnings.warn(
            f'overwriting query dependent operator "{func.__name__}" definition'
        )

    _QUERY_DEPENDENT_OP_REGISTRY[func_hash] = query_dependent_op(
        func, query_keys,
    )
    return func

def _intake_esm_get_keys_info(cat):
    """return a dictionary with the values of components of the keys in an
       intake catalog
    """
    groupby_attrs_values = {
        k: cat.unique(columns=k)[k]['values']
        for k in cat.groupby_attrs
    }
    key_info = {k: {} for k in cat.keys()}
    for key in cat.keys():
        for attr in cat.groupby_attrs:
            values = groupby_attrs_values[attr]
            match = [v for v in values if v in key]
            assert len(match) == 1, 'expecting a unique match'
            key_info[key][attr] = match[0]
    return key_info


#!/usr/bin/env python3

import psycopg2
from datetime import datetime

def get_data(cur, sql):

    cur.execute(sql)
    result = cur.fetchall()
    return (result)

def get_three_most_popular_articles(cur):

    """
    Input parameter: cursor of the database session
    Return parameter: None

    Query 1:
    What are the most popular three articles of all time
    """
    sql = "with log_detail as (select substr(path,10) path,id from log"\
        + " where length(path)>5)"   \
        + " ,art_det as (select a.title,det.id"\
        + " from log_detail det left join articles a "\
        + "on det.path = a.slug)"\
        + " select title,count(id) as access_count from "\
        + "art_det group by title"\
        + " order by 2 desc limit 3"
    res = get_data(cur, sql)
    print("The result of the most popular three articles of all time:\n")
    for rw in res:
        print("\"{0}\" --- {1} views".format(rw[0], rw[1]))

def get_most_popular_author(cur):

    """
    Input parameter: cursor of the database session
    Return parameter: None

    Query 2:
    What are the most popular articles author of all time
    """
    sql = "select a.name ,a.count_views from ("\
        + "select name ,count(id) as count_views "\
        + "from v_author_articles_access"\
        + " group by name order by 2 desc) a"
    res = get_data(cur, sql)
    print("The result of the most popular articles author of all time:\n")
    for rw in res:
        print("{} --- {} views".format(rw[0], rw[1]))

def get_days_with_onepc_failed(cur):

    """
    Input parameter: cursor of the database session
    Return parameter: None

    Query 3:
    On which days did more than 1% of requests lead to errors?
    """
    sql = "with day_rec as (select date_trunc('day', time) "\
        + "as days, status,id from log)"\
        + ",total as (select days,count(id) as total from "\
        + "day_rec group by days),"\
        + "error_total as (select days,count(id) as "\
        + "err_total from day_rec where status ="\
        + "'404 NOT FOUND' group by days ),"\
        + "err_percent as ("\
        + "select a.days,a.total,b.err_total,"\
        + "((b.err_total*1.00)/(a.total*1.00))*100.0 as "\
        + "percent from "\
        + "total a join error_total b on a.days = b.DAYS)"\
        + " select days ,round(percent,1) as err_txt "\
        + "from err_percent where percent>1"
    res = get_data(cur, sql)
    print("On which days did more than 1% of requests lead to errors:\n")
    for rw in res:
        print("{} --- {}% errors".format(
            datetime.strftime(rw[0], '%B %d, %Y'), rw[1]))

def main():

    conn = psycopg2.connect("dbname=news")
    cur = conn.cursor()
    get_three_most_popular_articles(cur)
    print("\n")
    get_most_popular_author(cur)
    print("\n")
    get_days_with_onepc_failed(cur)
    conn.close()

if __name__ == "__main__":
    main()


import torch
import torch.nn as nn
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = torch.device("cuda:0")

com = [
    16,
    "ELU",
    32,
    "ELU",
    64,
    "ELU",
    128,
    "ELU",
    256,
    "ELU",
    128,
    "ELU",
    64,
    "ELU",
    32,
    "ELU",
    4,
]
rec = [
    "SIGMOID",
    32,
    "ELU",
    64,
    "ELU",
    128,
    "ELU",
    256,
    "ELU",
    128,
    "ELU",
    64,
    "ELU",
    32,
    "ELU",
    16,
    "ELU",
    3,
]

class ComDefend(nn.Module):
    def __init__(self, phi=20):
        super(ComDefend, self).__init__()
        self.name = "ComDefend"
        self.phi = phi
        self.com, self.rec = self._make_layers(com, rec)

    def _make_layers(self, com, rec):
        comCNN = []
        recCNN = []
        inchannels = 3

        for x in com:
            if x == "ELU":
                comCNN.append(nn.ELU())
            elif x == "SIGMOID":
                comCNN.append(nn.Sigmoid())
            else:
                comCNN.append(nn.Conv2d(inchannels, x, kernel_size=3, padding=1))
                comCNN.append(nn.BatchNorm2d(x))
                inchannels = x

        for x in rec:
            if x == "ELU":
                recCNN.append(nn.ELU())
            elif x == "SIGMOID":
                recCNN.append(nn.Sigmoid())
            else:
                recCNN.append(nn.Conv2d(inchannels, x, kernel_size=3, padding=1))
                recCNN.append(nn.BatchNorm2d(x))
                inchannels = x

        return nn.Sequential(*comCNN), nn.Sequential(*recCNN)

    def forward(self, x):
        hidden = self.com(x)
        hidden_gauss = hidden + (self.phi ** 0.5) * torch.randn(hidden.shape).to(device)
        output = self.rec(hidden_gauss)

        return hidden, output

comDefend = ComDefend().to(device)

import numpy as np

data = (
    np.load("/home/liyanni/1307/zwh/whiteBox/data/cifar10_data.npy").astype("float32")
    / 255
)
data = data.transpose((0, 3, 1, 2))
label = np.load("/home/liyanni/1307/zwh/whiteBox/data/cifar10_label.npy")
label = np.reshape(label, (label.shape[0],))

def dataset(batch=50):
    train = data[:50000]
    for i in range(len(train) // batch):
        yield train[i * batch : (i + 1) * batch], train[i * batch : (i + 1) * batch]

comDefend.train()
from torch.autograd import Variable

criterion1 = nn.MSELoss()
criterion2 = nn.MSELoss()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, comDefend.parameters()))
for epoch in range(20):
    dataGen = dataset()
    l1 = 0.0
    l2 = 0.0
    for i, (x, y) in enumerate(dataGen, 1):
        # get the inputs
        x, y = torch.from_numpy(x), torch.from_numpy(y)
        x, y = x.to(device), y.to(device)
        x, y = Variable(x, requires_grad=True), Variable(y)
        # zero the parameter gradients
        optimizer.zero_grad()
        # forward
        o1, o2 = comDefend(x)
        # loss
        loss1 = criterion1(o1, torch.zeros(o1.shape).to(device))
        loss2 = criterion2(o2, y)
        loss = 0.0001 * loss1 + loss2
        # backward
        loss.backward()
        # update weights
        optimizer.step()
        # print statistics
        l1 += loss1.item()
        l2 += loss2.item()

        if i % 50 == 0:  # print every 2000 mini-batches
            print("epoch:{}  loss1:{}  loss2{}".format(epoch, l1 / 50.0, l2 / 50.0))
            l1, l2 = 0.0, 0.0

from models.resnet18 import ResNet50
from collections import OrderedDict

target = ResNet50()
pthfile = r"checkpoints/resnet50_ckpt.pth"
d = torch.load(pthfile)["net"]
d = OrderedDict([(k[7:], v) for (k, v) in d.items()])
target.load_state_dict(d)
target.to(device)

mean = np.mean(data, axis=0)

adv_x = (
    np.load("../adv_data/cifar10/cw/0.015x_adv.npy")
    .astype("float32")
    .transpose(0, 3, 1, 2)
)
adv_x = adv_x + mean

def testGen(data, batch=50):
    testx = data[50000:]
    testy = label[50000:]
    for i in range(len(testx) // batch):
        yield testx[i * batch : (i + 1) * batch], testy[i * batch : (i + 1) * batch]

test = testGen(adv_x)
correct = 0
total = 0
for inputs, labels in test:
    inputs, labels = torch.from_numpy(inputs), torch.from_numpy(labels)
    inputs, labels = inputs.to(device), labels.to(device)
    _, outputs = comDefend(inputs)
    outputs = target(outputs)
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum()
print("Accuracy {}".format(100.0 * correct / total))


import sys
import argparse
from os import mkdir, getpid
from os.path import isfile, isdir
import logging
import json
from datetime import date

if not isdir("./logs"):
    try:
        mkdir("./logs")
    except Exception as e:
        print("Failed to make log dir")

logfile = './logs/TrapHandiCap.log'

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s | %(levelname)s | %(funcName)s | %(module)s | %(message)s')

file_handler = logging.FileHandler(logfile)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(formatter)

logger.addHandler(file_handler)

logger.info("Logging started...")

class Calculator:
    """def __init__(self) -> None:
        self.number_of_scores = args["scores"]
        self.max_average = 24.00
        self.average = None
        self.handicap = None
        #self.data_file = self.init_data
        if not isfile(args["file"]):
            open(args["file"], 'x')
        self.data_file = args["file"]
        self.scores = self.load_data(self.data_file)"""
    def __init__(self) -> None:
        self.max_average = 24.00
        self.average = None
        self.handicap = None

    def __str__(self) -> str:
        if self.average != None or self.handicap != None:
            return f'Calculations({self.average}, {self.handicap})'
        else:
            return "No Calculation"

    def load_data(self, data_file) -> dict:
        try:
            file = open(data_file, 'r')
        except OSError as e:
            #logger.exception("Failed to read file " + str(data_file))
            print("Failed to read file " + str(data_file))
            sys.exit()
        data = json.load(file)
        #logger.info("Data loaded...")
        return data

    def sort_data(self) -> None:
        pass

    def calculate_average(self, scores: tuple) -> None:
        if len(scores) < self.number_of_scores:
            # If there are less than
            self.average = 0
            #logger.debug("The number of scores to calculate the average was less than number of scores to be used to calculate")
        else:
            # Calculate the average
            self.average = sum(scores) / len(scores)

    def calculate_handicap(self) -> None:
        if self.average > self.max_average:
            self.handicap = 0
        else:
            h = self.max_average - self.average
            self.handicap = round(h)

    def write_to_file(self, new_data: dict) -> None:
        with open(self.data_file, 'r+') as file:
            file_data = json.load(file)
            file_data['Scores'].append(new_data)
            file.seek(0)
            json.dump(file_data, file, indent=4)


from fakeredis import FakeStrictRedis

#from .testutils import BaseWRTests
from .testutils import FullStackTests

from pywb.warcserver.test.testutils import LiveServerTests

import os

import webtest

from six.moves.urllib.parse import quote, urlsplit
import time

general_req_data = "\
GET {path} HTTP/1.1\r\n\
Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8\r\n\
User-agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.116 Safari/537.36\r\n\
Host: {host}\r\n\
\r\n"

# ============================================================================
class TestWebRecRecorder(FullStackTests):
    @classmethod
    def setup_class(cls):
        super(TestWebRecRecorder, cls).setup_class(rec=False, no_app=True)

        from webrecorder.rec.main import init

        app = init()
        cls.wr_rec = app.wr
        cls.testapp = webtest.TestApp(app)

    @classmethod
    def teardown_class(cls):
        cls.wr_rec.writer.close()
        super(TestWebRecRecorder, cls).teardown_class()

    def _test_warc_write(self, url, user, coll, rec):
        parts = urlsplit(url)
        host = parts.netloc
        path = parts.path
        if parts.query:
            path += '?' + parts.query

        # Rec must be open
        #self.redis.hset('r:{user}:{coll}:{rec}:info'.format(user=user, coll=coll, rec=rec), 'id', rec)
        self.redis.setex('r:{rec}:open'.format(user=user, coll=coll, rec=rec), 30, 1)
        self.redis.hset('u:{user}:info'.format(user=user), 'size', 0)
        self.redis.hset('u:{user}:info'.format(user=user), 'max_size', 10000)

        req_url = '/record/live/resource/postreq?url={url}&param.recorder.user={user}&param.recorder.coll={coll}&param.recorder.rec={rec}'
        req_url = req_url.format(url=quote(url), user=user, coll=coll, rec=rec)
        resp = self.testapp.post(req_url, general_req_data.format(host=host, path=path).encode('utf-8'))

        while not self.wr_rec.recorder.write_queue.empty():
            self.wr_rec.recorder._write_one()

        assert resp.headers['Warcserver-Source-Coll'] == 'live'

        return resp

    def test_multi_user_rec(self):
        resp = self._test_warc_write('http://httpbin.org/get?foo=bar', user='USER', coll='COLL', rec='REC')

        keys = self.redis.keys()

        assert set(keys) == set([
            'r:REC:wk',
            'r:REC:cdxj',
            'r:REC:info',
            'r:REC:open',
            'r:REC:_ps',
            'r:REC:_pc',
            'c:COLL:info',
            'c:COLL:warc',
            'u:USER:info'
        ])

        resp.charset = 'utf-8'
        assert '"foo": "bar"' in resp.text

        size = self.redis.hget('r:REC:info', 'size')
        assert size is not None
        assert size == self.redis.hget('c:COLL:info', 'size')
        assert size == self.redis.hget('u:USER:info', 'size')

        assert self.redis.hget('r:REC:info', 'updated_at') is not None

    def test_multi_user_rec_2(self):
        resp = self._test_warc_write('http://httpbin.org/get?boo=far', user='USER', coll='COLL', rec='REC2')

        keys = self.redis.keys()

        assert set(keys) == set([
            'r:REC:wk',
            'r:REC:cdxj',
            'r:REC:info',
            'r:REC:open',
            'r:REC:_ps',
            'r:REC:_pc',
            'r:REC2:wk',
            'r:REC2:cdxj',
            'r:REC2:info',
            'r:REC2:open',
            'r:REC2:_ps',
            'r:REC2:_pc',
            'c:COLL:info',
            'c:COLL:warc',
            'u:USER:info'
        ])

        resp.charset = 'utf-8'
        assert '"boo": "far"' in resp.text


import esphome.codegen as cg
import esphome.config_validation as cv
from esphome import pins
from esphome.const import CONF_ID, CONF_PIN, \
    CONF_RESOLUTION, CONF_UNIT_OF_MEASUREMENT, UNIT_CELSIUS, \
    CONF_ICON, ICON_THERMOMETER, CONF_ACCURACY_DECIMALS

MULTI_CONF = True
AUTO_LOAD = ['sensor']

CONF_ONE_WIRE_ID = 'one_wire_id'
CONF_AUTO_SETUP_SENSORS = 'auto_setup_sensors'
CONF_SENSOR_NAME_TEMPLATE = 'sensor_name_template'
SENSOR_NAME_TEMPLATE_DEFAULT = '%s.%s'

dallas_ns = cg.esphome_ns.namespace('dallas')
DallasComponent = dallas_ns.class_('DallasComponent', cg.PollingComponent)
ESPOneWire = dallas_ns.class_('ESPOneWire')

CONFIG_SCHEMA = cv.Schema({
    cv.GenerateID(): cv.declare_id(DallasComponent),
    cv.GenerateID(CONF_ONE_WIRE_ID): cv.declare_id(ESPOneWire),
    cv.Required(CONF_PIN): pins.gpio_input_pin_schema,
    cv.Optional(CONF_AUTO_SETUP_SENSORS, default=False): cv.boolean,
    cv.Optional(CONF_SENSOR_NAME_TEMPLATE, default=SENSOR_NAME_TEMPLATE_DEFAULT): cv.string_strict,
    cv.Optional(CONF_RESOLUTION, default=12): cv.int_range(min=9, max=12),
    cv.Optional(CONF_UNIT_OF_MEASUREMENT, default=UNIT_CELSIUS): cv.string_strict,
    cv.Optional(CONF_ICON, default=ICON_THERMOMETER): cv.icon,
    cv.Optional(CONF_ACCURACY_DECIMALS, default=1): cv.int_,
}).extend(cv.polling_component_schema('60s'))

def to_code(config):
    pin = yield cg.gpio_pin_expression(config[CONF_PIN])
    one_wire = cg.new_Pvariable(config[CONF_ONE_WIRE_ID], pin)
    var = cg.new_Pvariable(config[CONF_ID], one_wire)
    if CONF_AUTO_SETUP_SENSORS in config:
        cg.add(var.set_auto_setup_sensors(config[CONF_AUTO_SETUP_SENSORS]))
    if CONF_SENSOR_NAME_TEMPLATE in config:
        cg.add(var.set_sensor_name_template(config[CONF_SENSOR_NAME_TEMPLATE]))
    if CONF_RESOLUTION in config:
        cg.add(var.set_resolution(config[CONF_RESOLUTION]))
    if CONF_UNIT_OF_MEASUREMENT in config:
        cg.add(var.set_unit_of_measurement(config[CONF_UNIT_OF_MEASUREMENT]))
    if CONF_ICON in config:
        cg.add(var.set_icon(config[CONF_ICON]))
    if CONF_ACCURACY_DECIMALS in config:
        cg.add(var.set_accuracy_decimals(config[CONF_ACCURACY_DECIMALS]))
    yield cg.register_component(var, config)


import pytest
import requests

from quark.utils.output import output_parent_function_table
from quark.Objects.quark import Quark

APK_SOURCE = (
    "https://github.com/quark-engine/apk-malware-samples"
    "/raw/master/14d9f1a92dd984d6040cc41ed06e273e.apk"
)
APK_FILENAME = "14d9f1a92dd984d6040cc41ed06e273e.apk"

@pytest.fixture()
def quark_obj(scope="function"):
    r = requests.get(APK_SOURCE, allow_redirects=True)
    open(APK_FILENAME, "wb").write(r.content)

    return Quark(APK_FILENAME)

def test_output_parent_function_table(capsys, quark_obj):
    method_obj = quark_obj.apkinfo.find_method(
        "Lcom/google/progress/WifiCheckTask;", "goConnectNetwork", "()V"
    )

    call_graph_analysis_list = (
        {"crime": "The Crime", "parent": method_obj},
        {"crime": "The Crime", "parent": method_obj},
        {"crime": "Another Crime", "parent": method_obj},
    )

    # Test one crime
    output_parent_function_table([call_graph_analysis_list[0]])

    output = capsys.readouterr().out
    assert output.count("Lcom/google/progress/WifiCheckTask;goConnectNetwork") == 1
    assert output.count("The Crime") == 1

    # Test crimes with duplicated description
    output_parent_function_table(call_graph_analysis_list)

    output = capsys.readouterr().out
    print(output)
    assert output.count("Lcom/google/progress/WifiCheckTask;goConnectNetwork") == 1
    assert output.count("The Crime") == 1
    assert output.count("Another Crime") == 1