<commit_before>#!/usr/bin/env python3

import os
import time
import datetime
from subprocess import call
import shutil

import terminal_colors
import path_utils
import pakgen
import encrypt
import dirsize
import sha256_wrapper
import shred_wrapper
import create_and_write_file

SCRIPT_FOLDER = os.path.dirname(os.path.realpath(__file__))

class BackupEngine:

    def __init__(_self, bkprep, bkarts, bktgs_root, bktgs_base, bktmp, bkwarns, pphrase):
        _self.BKPREPARATION = bkprep
        _self.BKARTIFACTS = bkarts
        _self.BKTARGETS_ROOT = bktgs_root
        _self.BKTARGETS_BASEDIR = bktgs_base
        _self.BKTEMP = bktmp
        _self.BKWARNINGS = bkwarns
        _self.PASSPHRASE = pphrase

    def gettimestamp(_self):
        return datetime.datetime.fromtimestamp(time.time()).strftime('%H:%M:%S %d-%m-%Y')

    def run(_self):

        # validations
        art_filtered = []
        for it in _self.BKARTIFACTS:
            if not os.path.exists(it[0]):
                if it[1]:
                    print("%sThe path [%s] is marked for backing up, but does not exist. Aborting%s" % (terminal_colors.TTY_RED, it[0], terminal_colors.TTY_WHITE))
                    return False
                else:
                    print("%sThe path [%s] is marked for backing up, but does not exist.%s" % (terminal_colors.TTY_YELLOW_BOLD, it[0], terminal_colors.TTY_WHITE))
            else:
                art_filtered.append(it)
        _self.BKARTIFACTS = art_filtered

        if os.path.exists(_self.BKTEMP):
            print("%s[%s] already exists. For safety reasons, this script is aborted.%s" % (terminal_colors.TTY_RED, _self.BKTEMP, terminal_colors.TTY_WHITE))
            return False

        print("%sBeginning backup operations at %s.%s" % (terminal_colors.TTY_GREEN, _self.gettimestamp(), terminal_colors.TTY_WHITE))

        for it in _self.BKTARGETS_ROOT:
            if not os.path.isdir(it):
                print("%sThe path [%s] is marked as a writing target, but does not exist. Aborting.%s" % (terminal_colors.TTY_RED, it, terminal_colors.TTY_WHITE))
                return False

        if len(_self.BKPREPARATION) == 2:
            if len(_self.BKPREPARATION[0]) > 0:
                print("%sPreparing...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))
                prep_cmd = [_self.BKPREPARATION[0]]
                for prep_arg in _self.BKPREPARATION[1]:
                    prep_cmd.append(prep_arg)
                prep_r = call(prep_cmd)
                if prep_r != 0:
                    print("%sFailed preparing backup. Aborting.%s" % (terminal_colors.TTY_RED, terminal_colors.TTY_WHITE))
                    return False

        print("%sDeleting old backup...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))
        for it in _self.BKTARGETS_ROOT:
            test_subj = path_utils.concat_path(it, _self.BKTARGETS_BASEDIR)
            if not path_utils.scratchfolder(test_subj):
                print("%sCannot scratch %s - are the external media available/attached?%s" % (terminal_colors.TTY_RED, it, terminal_colors.TTY_WHITE))
                return False
            else:
                shutil.rmtree(test_subj) # redundant but necessary

        print("%sCreating backup...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))

        if not path_utils.scratchfolder(_self.BKTEMP):
            print("%sUnable to create temporary path [%s]. Aborting.%s" % (terminal_colors.TTY_RED, _self.BKTEMP, terminal_colors.TTY_WHITE))
            return False
        BKTEMP_AND_BASEDIR = path_utils.concat_path(_self.BKTEMP, _self.BKTARGETS_BASEDIR)
        os.mkdir(BKTEMP_AND_BASEDIR)
        if not os.path.exists(BKTEMP_AND_BASEDIR):
            print("%sUnable to create temporary+base path [%s]. Aborting.%s" % (terminal_colors.TTY_RED, BKTEMP_AND_BASEDIR, terminal_colors.TTY_WHITE))
            return False
        with open(path_utils.concat_path(BKTEMP_AND_BASEDIR, "bk_date.txt"), "w+") as f:
            f.write(_self.gettimestamp() + "\n")

        for it in _self.BKARTIFACTS:
            print("%sCurrent: %s, started at %s%s" % (terminal_colors.TTY_BLUE, it[0], datetime.datetime.fromtimestamp(time.time()).strftime('%H:%M:%S'), terminal_colors.TTY_WHITE))
            BKTMP_PLUS_ARTBASE = path_utils.concat_path(BKTEMP_AND_BASEDIR, path_utils.basename_filtered(os.path.dirname(it[0])))
            path_utils.guaranteefolder(BKTMP_PLUS_ARTBASE)

            CURPAK = path_utils.concat_path(BKTMP_PLUS_ARTBASE, path_utils.basename_filtered(it[0]))
            CURPAK_TAR_BZ2 = CURPAK + ".tar.bz2"
            CURPAK_TAR_BZ2_ENC = CURPAK_TAR_BZ2 + ".enc"
            CURPAK_TAR_BZ2_ENC_HASH = CURPAK_TAR_BZ2_ENC + ".sha256"

            # create the package
            v, r = pakgen.pakgen(CURPAK, False, [it[0]]) # hash will be generated later (from the encrypted package)
            if not v:
                print("%sFailed generating [%s].%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2, terminal_colors.TTY_WHITE))
                return False
            if len(r) > 0:
                print("Output from pakgen: %s" % r)

            # encrypt plain package
            v, r = encrypt.symmetric_encrypt(CURPAK_TAR_BZ2, CURPAK_TAR_BZ2_ENC, _self.PASSPHRASE)
            if not v:
                print("%sFailed encrypting package: [%s].%s" % (terminal_colors.TTY_RED, r, terminal_colors.TTY_WHITE))
                return False

            # shred plain package
            v, r = shred_wrapper.shred_target(CURPAK_TAR_BZ2)
            if not v:
                print("%sFailed shredding plain package: [%s].%s" % (terminal_colors.TTY_RED, r, terminal_colors.TTY_WHITE))
                return False

            warn_size_each_active_local = (_self.BKWARNINGS[0][0] is not None) and (_self.BKWARNINGS[0][1] is not None)
            warn_size_each_local = _self.BKWARNINGS[0][0]
            warn_size_each_abort_local = _self.BKWARNINGS[0][1]

            if it[2] is not None:
                # warning size option overridden
                warn_size_each_active_local = True
                warn_size_each_local = it[2]

            if it[3] is not None:
                # warning abort option overridden
                warn_size_each_active_local = True
                warn_size_each_abort_local = it[3]

            if warn_size_each_active_local:
                if dirsize.get_dir_size(CURPAK_TAR_BZ2_ENC, False) > warn_size_each_local:
                    if warn_size_each_abort_local:
                        print("%sGenerated package [%s] exceeds the size limit. Aborting.%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2_ENC, terminal_colors.TTY_WHITE))
                        return False
                    else:
                        print("%sGenerated package [%s] exceeds the size limit.%s" % (terminal_colors.TTY_YELLOW_BOLD, CURPAK_TAR_BZ2_ENC, terminal_colors.TTY_WHITE))

            # create hash from the encrypted package
            v, r = sha256_wrapper.hash_sha_256_app_file(CURPAK_TAR_BZ2_ENC)
            if not v:
                print("%sFailed generating hash for [%s].%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2_ENC, terminal_colors.TTY_WHITE))
                return False
            create_and_write_file.create_file_contents(CURPAK_TAR_BZ2_ENC_HASH, r)

        warn_size_final_active_local = (_self.BKWARNINGS[1][0] is not None) and (_self.BKWARNINGS[1][1] is not None)
        if warn_size_final_active_local:
            if dirsize.get_dir_size(BKTEMP_AND_BASEDIR, False) > _self.BKWARNINGS[1][0]:
                if _self.BKWARNINGS[1][1]: # abort
                    print("%sGenerated backup exceeds size limit. Aborting.%s" % (terminal_colors.TTY_RED, terminal_colors.TTY_WHITE))
                    return False
                else:
                    print("%sGenerated backup exceeds size limit.%s" % (terminal_colors.TTY_YELLOW_BOLD, terminal_colors.TTY_WHITE))

        print("%sWriting to targets...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))

        for it in _self.BKTARGETS_ROOT:
            shutil.copytree(BKTEMP_AND_BASEDIR, path_utils.concat_path(it, _self.BKTARGETS_BASEDIR))
            call(["umount", it])

        shutil.rmtree(_self.BKTEMP)
        print("%sDone at %s%s" % (terminal_colors.TTY_GREEN, _self.gettimestamp(), terminal_colors.TTY_WHITE))

        return True
<commit_msg>patching backup_engine to prevent overwrites<commit_after>#!/usr/bin/env python3

import os
import time
import datetime
from subprocess import call
import shutil

import terminal_colors
import path_utils
import pakgen
import encrypt
import dirsize
import sha256_wrapper
import shred_wrapper
import create_and_write_file

SCRIPT_FOLDER = os.path.dirname(os.path.realpath(__file__))

class BackupEngine:

    def __init__(_self, bkprep, bkarts, bktgs_root, bktgs_base, bktmp, bkwarns, pphrase):
        _self.BKPREPARATION = bkprep
        _self.BKARTIFACTS = bkarts
        _self.BKTARGETS_ROOT = bktgs_root
        _self.BKTARGETS_BASEDIR = bktgs_base
        _self.BKTEMP = bktmp
        _self.BKWARNINGS = bkwarns
        _self.PASSPHRASE = pphrase

    def gettimestamp(_self):
        return datetime.datetime.fromtimestamp(time.time()).strftime('%H:%M:%S %d-%m-%Y')

    def run(_self):

        # validations
        art_filtered = []
        for it in _self.BKARTIFACTS:
            if not os.path.exists(it[0]):
                if it[1]:
                    print("%sThe path [%s] is marked for backing up, but does not exist. Aborting%s" % (terminal_colors.TTY_RED, it[0], terminal_colors.TTY_WHITE))
                    return False
                else:
                    print("%sThe path [%s] is marked for backing up, but does not exist.%s" % (terminal_colors.TTY_YELLOW_BOLD, it[0], terminal_colors.TTY_WHITE))
            else:
                art_filtered.append(it)
        _self.BKARTIFACTS = art_filtered

        if os.path.exists(_self.BKTEMP):
            print("%s[%s] already exists. For safety reasons, this script is aborted.%s" % (terminal_colors.TTY_RED, _self.BKTEMP, terminal_colors.TTY_WHITE))
            return False

        print("%sBeginning backup operations at %s.%s" % (terminal_colors.TTY_GREEN, _self.gettimestamp(), terminal_colors.TTY_WHITE))

        for it in _self.BKTARGETS_ROOT:
            if not os.path.isdir(it):
                print("%sThe path [%s] is marked as a writing target, but does not exist. Aborting.%s" % (terminal_colors.TTY_RED, it, terminal_colors.TTY_WHITE))
                return False

        if len(_self.BKPREPARATION) == 2:
            if len(_self.BKPREPARATION[0]) > 0:
                print("%sPreparing...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))
                prep_cmd = [_self.BKPREPARATION[0]]
                for prep_arg in _self.BKPREPARATION[1]:
                    prep_cmd.append(prep_arg)
                prep_r = call(prep_cmd)
                if prep_r != 0:
                    print("%sFailed preparing backup. Aborting.%s" % (terminal_colors.TTY_RED, terminal_colors.TTY_WHITE))
                    return False

        print("%sDeleting old backup...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))
        for it in _self.BKTARGETS_ROOT:
            test_subj = path_utils.concat_path(it, _self.BKTARGETS_BASEDIR)
            if not path_utils.scratchfolder(test_subj):
                print("%sCannot scratch %s - are the external media available/attached?%s" % (terminal_colors.TTY_RED, it, terminal_colors.TTY_WHITE))
                return False
            else:
                shutil.rmtree(test_subj) # redundant but necessary

        print("%sCreating backup...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))

        if not path_utils.scratchfolder(_self.BKTEMP):
            print("%sUnable to create temporary path [%s]. Aborting.%s" % (terminal_colors.TTY_RED, _self.BKTEMP, terminal_colors.TTY_WHITE))
            return False
        BKTEMP_AND_BASEDIR = path_utils.concat_path(_self.BKTEMP, _self.BKTARGETS_BASEDIR)
        os.mkdir(BKTEMP_AND_BASEDIR)
        if not os.path.exists(BKTEMP_AND_BASEDIR):
            print("%sUnable to create temporary+base path [%s]. Aborting.%s" % (terminal_colors.TTY_RED, BKTEMP_AND_BASEDIR, terminal_colors.TTY_WHITE))
            return False
        with open(path_utils.concat_path(BKTEMP_AND_BASEDIR, "bk_date.txt"), "w+") as f:
            f.write(_self.gettimestamp() + "\n")

        for it in _self.BKARTIFACTS:
            print("%sCurrent: %s, started at %s%s" % (terminal_colors.TTY_BLUE, it[0], datetime.datetime.fromtimestamp(time.time()).strftime('%H:%M:%S'), terminal_colors.TTY_WHITE))
            BKTMP_PLUS_ARTBASE = path_utils.concat_path(BKTEMP_AND_BASEDIR, path_utils.basename_filtered(os.path.dirname(it[0])))
            if path_utils.basename_filtered(BKTMP_PLUS_ARTBASE) == path_utils.basename_filtered(BKTEMP_AND_BASEDIR):
                BKTMP_PLUS_ARTBASE = path_utils.concat_path(BKTMP_PLUS_ARTBASE, "(root)")
            path_utils.guaranteefolder(BKTMP_PLUS_ARTBASE)

            CURPAK = path_utils.concat_path(BKTMP_PLUS_ARTBASE, path_utils.basename_filtered(it[0]))
            CURPAK_TAR_BZ2 = CURPAK + ".tar.bz2"
            CURPAK_TAR_BZ2_ENC = CURPAK_TAR_BZ2 + ".enc"
            CURPAK_TAR_BZ2_ENC_HASH = CURPAK_TAR_BZ2_ENC + ".sha256"

            if os.path.exists(CURPAK) or os.path.exists(CURPAK_TAR_BZ2) or os.path.exists(CURPAK_TAR_BZ2_ENC) or os.path.exists(CURPAK_TAR_BZ2_ENC_HASH):
                print("%sFailed generating [%s] - duplicated artifact - aborting to avoid overwrites.%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2, terminal_colors.TTY_WHITE))
                return False

            # create the package
            v, r = pakgen.pakgen(CURPAK, False, [it[0]]) # hash will be generated later (from the encrypted package)
            if not v:
                print("%sFailed generating [%s].%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2, terminal_colors.TTY_WHITE))
                return False
            if len(r) > 0:
                print("Output from pakgen: %s" % r)

            # encrypt plain package
            v, r = encrypt.symmetric_encrypt(CURPAK_TAR_BZ2, CURPAK_TAR_BZ2_ENC, _self.PASSPHRASE)
            if not v:
                print("%sFailed encrypting package: [%s].%s" % (terminal_colors.TTY_RED, r, terminal_colors.TTY_WHITE))
                return False

            # shred plain package
            v, r = shred_wrapper.shred_target(CURPAK_TAR_BZ2)
            if not v:
                print("%sFailed shredding plain package: [%s].%s" % (terminal_colors.TTY_RED, r, terminal_colors.TTY_WHITE))
                return False

            warn_size_each_active_local = (_self.BKWARNINGS[0][0] is not None) and (_self.BKWARNINGS[0][1] is not None)
            warn_size_each_local = _self.BKWARNINGS[0][0]
            warn_size_each_abort_local = _self.BKWARNINGS[0][1]

            if it[2] is not None:
                # warning size option overridden
                warn_size_each_active_local = True
                warn_size_each_local = it[2]

            if it[3] is not None:
                # warning abort option overridden
                warn_size_each_active_local = True
                warn_size_each_abort_local = it[3]

            if warn_size_each_active_local:
                if dirsize.get_dir_size(CURPAK_TAR_BZ2_ENC, False) > warn_size_each_local:
                    if warn_size_each_abort_local:
                        print("%sGenerated package [%s] exceeds the size limit. Aborting.%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2_ENC, terminal_colors.TTY_WHITE))
                        return False
                    else:
                        print("%sGenerated package [%s] exceeds the size limit.%s" % (terminal_colors.TTY_YELLOW_BOLD, CURPAK_TAR_BZ2_ENC, terminal_colors.TTY_WHITE))

            # create hash from the encrypted package
            v, r = sha256_wrapper.hash_sha_256_app_file(CURPAK_TAR_BZ2_ENC)
            if not v:
                print("%sFailed generating hash for [%s].%s" % (terminal_colors.TTY_RED, CURPAK_TAR_BZ2_ENC, terminal_colors.TTY_WHITE))
                return False
            create_and_write_file.create_file_contents(CURPAK_TAR_BZ2_ENC_HASH, r)

        warn_size_final_active_local = (_self.BKWARNINGS[1][0] is not None) and (_self.BKWARNINGS[1][1] is not None)
        if warn_size_final_active_local:
            if dirsize.get_dir_size(BKTEMP_AND_BASEDIR, False) > _self.BKWARNINGS[1][0]:
                if _self.BKWARNINGS[1][1]: # abort
                    print("%sGenerated backup exceeds size limit. Aborting.%s" % (terminal_colors.TTY_RED, terminal_colors.TTY_WHITE))
                    return False
                else:
                    print("%sGenerated backup exceeds size limit.%s" % (terminal_colors.TTY_YELLOW_BOLD, terminal_colors.TTY_WHITE))

        print("%sWriting to targets...%s" % (terminal_colors.TTY_BLUE, terminal_colors.TTY_WHITE))

        for it in _self.BKTARGETS_ROOT:
            shutil.copytree(BKTEMP_AND_BASEDIR, path_utils.concat_path(it, _self.BKTARGETS_BASEDIR))
            call(["umount", it])

        shutil.rmtree(_self.BKTEMP)
        print("%sDone at %s%s" % (terminal_colors.TTY_GREEN, _self.gettimestamp(), terminal_colors.TTY_WHITE))

        return True
<|endoftext|>


<commit_before>import unittest

import crankshaft.pysal_utils as pu
from crankshaft import random_seeds
from collections import OrderedDict

class PysalUtilsTest(unittest.TestCase):
    """Testing class for utility functions related to PySAL integrations"""

    def setUp(self):
        self.params = OrderedDict([("id_col", "cartodb_id"),
                                   ("attr1", "andy"),
                                   ("attr2", "jay_z"),
                                   ("subquery", "SELECT * FROM a_list"),
                                   ("geom_col", "the_geom"),
                                   ("num_ngbrs", 321)])

        self.params_array = {"id_col": "cartodb_id",
                             "time_cols": ["_2013_dec", "_2014_jan", "_2014_feb"],
                             "subquery": "SELECT * FROM a_list",
                             "geom_col": "the_geom",
                             "num_ngbrs": 321}

    def test_query_attr_select(self):
        """Test query_attr_select"""

        ans = "i.\"andy\"::numeric As attr1, " \
              "i.\"jay_z\"::numeric As attr2, "

        ans_array = "i.\"_2013_dec\"::numeric As attr1, " \
                    "i.\"_2014_jan\"::numeric As attr2, " \
                    "i.\"_2014_feb\"::numeric As attr3, "

        self.assertEqual(pu.query_attr_select(self.params), ans)
        self.assertEqual(pu.query_attr_select(self.params_array), ans_array)

    def test_query_attr_where(self):
        """Test pu.query_attr_where"""

        ans = "idx_replace.\"andy\" IS NOT NULL AND " \
              "idx_replace.\"jay_z\" IS NOT NULL AND " \
              "idx_replace.\"jay_z\" <> 0"

        ans_array = "idx_replace.\"_2013_dec\" IS NOT NULL AND " \
                    "idx_replace.\"_2014_jan\" IS NOT NULL AND " \
                    "idx_replace.\"_2014_feb\" IS NOT NULL"

        self.assertEqual(pu.query_attr_where(self.params), ans)
        self.assertEqual(pu.query_attr_where(self.params_array), ans_array)

    def test_knn(self):
        """Test knn neighbors constructor"""

        ans = "SELECT i.\"cartodb_id\" As id, " \
                     "i.\"andy\"::numeric As attr1, " \
                     "i.\"jay_z\"::numeric As attr2, " \
                     "(SELECT ARRAY(SELECT j.\"cartodb_id\" " \
                                   "FROM (SELECT * FROM a_list) As j " \
                                   "WHERE " \
                                    "i.\"cartodb_id\" <> j.\"cartodb_id\" AND " \
                                    "j.\"andy\" IS NOT NULL AND " \
                                    "j.\"jay_z\" IS NOT NULL AND " \
                                    "j.\"jay_z\" <> 0 " \
                                   "ORDER BY " \
                                    "j.\"the_geom\" <-> i.\"the_geom\" ASC " \
                      "LIMIT 321)) As neighbors " \
              "FROM (SELECT * FROM a_list) As i " \
              "WHERE i.\"andy\" IS NOT NULL AND " \
                    "i.\"jay_z\" IS NOT NULL AND " \
                    "i.\"jay_z\" <> 0 " \
              "ORDER BY i.\"cartodb_id\" ASC;"

        ans_array = "SELECT i.\"cartodb_id\" As id, " \
              "i.\"_2013_dec\"::numeric As attr1, " \
              "i.\"_2014_jan\"::numeric As attr2, " \
              "i.\"_2014_feb\"::numeric As attr3, " \
              "(SELECT ARRAY(SELECT j.\"cartodb_id\" " \
                            "FROM (SELECT * FROM a_list) As j " \
                            "WHERE i.\"cartodb_id\" <> j.\"cartodb_id\" AND " \
                                  "j.\"_2013_dec\" IS NOT NULL AND " \
                                  "j.\"_2014_jan\" IS NOT NULL AND " \
                                  "j.\"_2014_feb\" IS NOT NULL " \
                            "ORDER BY j.\"the_geom\" <-> i.\"the_geom\" ASC " \
                            "LIMIT 321)) As neighbors " \
              "FROM (SELECT * FROM a_list) As i " \
              "WHERE i.\"_2013_dec\" IS NOT NULL AND " \
                    "i.\"_2014_jan\" IS NOT NULL AND " \
                    "i.\"_2014_feb\" IS NOT NULL "\
              "ORDER BY i.\"cartodb_id\" ASC;"

        self.assertEqual(pu.knn(self.params), ans)
        self.assertEqual(pu.knn(self.params_array), ans_array)

    def test_queen(self):
        """Test queen neighbors constructor"""

        ans = "SELECT i.\"cartodb_id\" As id, " \
                     "i.\"andy\"::numeric As attr1, " \
                     "i.\"jay_z\"::numeric As attr2, " \
                     "(SELECT ARRAY(SELECT j.\"cartodb_id\" " \
                                   "FROM (SELECT * FROM a_list) As j " \
                                   "WHERE " \
                                   "i.\"cartodb_id\" <> j.\"cartodb_id\" AND " \
                                   "ST_Touches(i.\"the_geom\", " \
                                              "j.\"the_geom\") AND " \
                                   "j.\"andy\" IS NOT NULL AND " \
                                   "j.\"jay_z\" IS NOT NULL AND " \
                                   "j.\"jay_z\" <> 0)" \
                                  ") As neighbors " \
              "FROM (SELECT * FROM a_list) As i " \
              "WHERE i.\"andy\" IS NOT NULL AND " \
                    "i.\"jay_z\" IS NOT NULL AND " \
                    "i.\"jay_z\" <> 0 " \
              "ORDER BY i.\"cartodb_id\" ASC;"

        self.assertEqual(pu.queen(self.params), ans)

    def test_construct_neighbor_query(self):
        """Test construct_neighbor_query"""

        # Compare to raw knn query
        self.assertEqual(pu.construct_neighbor_query('knn', self.params),
                         pu.knn(self.params))

    def test_get_attributes(self):
        """Test get_attributes"""

        ## need to add tests

        self.assertEqual(True, True)

    def test_get_weight(self):
        """Test get_weight"""

        self.assertEqual(True, True)

    def test_empty_zipped_array(self):
        """Test empty_zipped_array"""
        ans2 = [(None, None)]
        ans4 = [(None, None, None, None)]
        self.assertEqual(pu.empty_zipped_array(2), ans2)
        self.assertEqual(pu.empty_zipped_array(4), ans4)
<commit_msg>adding more tests<commit_after>import unittest

import crankshaft.pysal_utils as pu
from crankshaft import random_seeds
from collections import OrderedDict

class PysalUtilsTest(unittest.TestCase):
    """Testing class for utility functions related to PySAL integrations"""

    def setUp(self):
        self.params = OrderedDict([("id_col", "cartodb_id"),
                                   ("attr1", "andy"),
                                   ("attr2", "jay_z"),
                                   ("subquery", "SELECT * FROM a_list"),
                                   ("geom_col", "the_geom"),
                                   ("num_ngbrs", 321)])

        self.params2 = OrderedDict([("id_col", "cartodb_id"),
                                    ("attr1", "price"),
                                    ("attr2", "sq_meters"),
                                    ("subquery", "SELECT * FROM a_list"),
                                    ("geom_col", "the_geom"),
                                    ("num_ngbrs", 321)])

        self.params_array = {"id_col": "cartodb_id",
                             "time_cols": ["_2013_dec", "_2014_jan", "_2014_feb"],
                             "subquery": "SELECT * FROM a_list",
                             "geom_col": "the_geom",
                             "num_ngbrs": 321}

    def test_query_attr_select(self):
        """Test query_attr_select"""

        ans = "i.\"andy\"::numeric As attr1, " \
              "i.\"jay_z\"::numeric As attr2, "

        ans2 = "i.\"price\"::numeric As attr1, " \
               "i.\"sq_meters\"::numeric As attr2, "

        ans_array = "i.\"_2013_dec\"::numeric As attr1, " \
                    "i.\"_2014_jan\"::numeric As attr2, " \
                    "i.\"_2014_feb\"::numeric As attr3, "

        self.assertEqual(pu.query_attr_select(self.params), ans)
        self.assertEqual(pu.query_attr_select(self.params2), ans2)
        self.assertEqual(pu.query_attr_select(self.params_array), ans_array)

    def test_query_attr_where(self):
        """Test pu.query_attr_where"""

        ans = "idx_replace.\"andy\" IS NOT NULL AND " \
              "idx_replace.\"jay_z\" IS NOT NULL AND " \
              "idx_replace.\"jay_z\" <> 0"

        ans_array = "idx_replace.\"_2013_dec\" IS NOT NULL AND " \
                    "idx_replace.\"_2014_jan\" IS NOT NULL AND " \
                    "idx_replace.\"_2014_feb\" IS NOT NULL"

        self.assertEqual(pu.query_attr_where(self.params), ans)
        self.assertEqual(pu.query_attr_where(self.params_array), ans_array)

    def test_knn(self):
        """Test knn neighbors constructor"""

        ans = "SELECT i.\"cartodb_id\" As id, " \
                     "i.\"andy\"::numeric As attr1, " \
                     "i.\"jay_z\"::numeric As attr2, " \
                     "(SELECT ARRAY(SELECT j.\"cartodb_id\" " \
                                   "FROM (SELECT * FROM a_list) As j " \
                                   "WHERE " \
                                    "i.\"cartodb_id\" <> j.\"cartodb_id\" AND " \
                                    "j.\"andy\" IS NOT NULL AND " \
                                    "j.\"jay_z\" IS NOT NULL AND " \
                                    "j.\"jay_z\" <> 0 " \
                                   "ORDER BY " \
                                    "j.\"the_geom\" <-> i.\"the_geom\" ASC " \
                      "LIMIT 321)) As neighbors " \
              "FROM (SELECT * FROM a_list) As i " \
              "WHERE i.\"andy\" IS NOT NULL AND " \
                    "i.\"jay_z\" IS NOT NULL AND " \
                    "i.\"jay_z\" <> 0 " \
              "ORDER BY i.\"cartodb_id\" ASC;"

        ans_array = "SELECT i.\"cartodb_id\" As id, " \
              "i.\"_2013_dec\"::numeric As attr1, " \
              "i.\"_2014_jan\"::numeric As attr2, " \
              "i.\"_2014_feb\"::numeric As attr3, " \
              "(SELECT ARRAY(SELECT j.\"cartodb_id\" " \
                            "FROM (SELECT * FROM a_list) As j " \
                            "WHERE i.\"cartodb_id\" <> j.\"cartodb_id\" AND " \
                                  "j.\"_2013_dec\" IS NOT NULL AND " \
                                  "j.\"_2014_jan\" IS NOT NULL AND " \
                                  "j.\"_2014_feb\" IS NOT NULL " \
                            "ORDER BY j.\"the_geom\" <-> i.\"the_geom\" ASC " \
                            "LIMIT 321)) As neighbors " \
              "FROM (SELECT * FROM a_list) As i " \
              "WHERE i.\"_2013_dec\" IS NOT NULL AND " \
                    "i.\"_2014_jan\" IS NOT NULL AND " \
                    "i.\"_2014_feb\" IS NOT NULL "\
              "ORDER BY i.\"cartodb_id\" ASC;"

        self.assertEqual(pu.knn(self.params), ans)
        self.assertEqual(pu.knn(self.params_array), ans_array)

    def test_queen(self):
        """Test queen neighbors constructor"""

        ans = "SELECT i.\"cartodb_id\" As id, " \
                     "i.\"andy\"::numeric As attr1, " \
                     "i.\"jay_z\"::numeric As attr2, " \
                     "(SELECT ARRAY(SELECT j.\"cartodb_id\" " \
                                   "FROM (SELECT * FROM a_list) As j " \
                                   "WHERE " \
                                   "i.\"cartodb_id\" <> j.\"cartodb_id\" AND " \
                                   "ST_Touches(i.\"the_geom\", " \
                                              "j.\"the_geom\") AND " \
                                   "j.\"andy\" IS NOT NULL AND " \
                                   "j.\"jay_z\" IS NOT NULL AND " \
                                   "j.\"jay_z\" <> 0)" \
                                  ") As neighbors " \
              "FROM (SELECT * FROM a_list) As i " \
              "WHERE i.\"andy\" IS NOT NULL AND " \
                    "i.\"jay_z\" IS NOT NULL AND " \
                    "i.\"jay_z\" <> 0 " \
              "ORDER BY i.\"cartodb_id\" ASC;"

        self.assertEqual(pu.queen(self.params), ans)

    def test_construct_neighbor_query(self):
        """Test construct_neighbor_query"""

        # Compare to raw knn query
        self.assertEqual(pu.construct_neighbor_query('knn', self.params),
                         pu.knn(self.params))

    def test_get_attributes(self):
        """Test get_attributes"""

        ## need to add tests

        self.assertEqual(True, True)

    def test_get_weight(self):
        """Test get_weight"""

        self.assertEqual(True, True)

    def test_empty_zipped_array(self):
        """Test empty_zipped_array"""
        ans2 = [(None, None)]
        ans4 = [(None, None, None, None)]
        self.assertEqual(pu.empty_zipped_array(2), ans2)
        self.assertEqual(pu.empty_zipped_array(4), ans4)
<|endoftext|>


<commit_before>import glob

from .exceptions import RestoreError
from .settings import BACKUP_DIR, BACKUP_TIMESTAMP_PATTERN

def get_latest_timestamped_file(ext):
    """
    Gets the latest timestamped backup file name with the given database type
    extension.
    """
    # Make glob pattern
    pattern = '{dir}/{pattern}.{ext}.gz'.format(
        dir=BACKUP_DIR,
        pattern=BACKUP_TIMESTAMP_PATTERN,
        ext=ext,
    )

    # Find files
    l = glob.glob(pattern)
    l.sort()
    l.reverse()

    if not l:
        raise RestoreError("No backups found matching '{0}' pattern".format(pattern))

    return l[0]
<commit_msg>Comments<commit_after>import glob

from .exceptions import RestoreError
from .settings import BACKUP_DIR, BACKUP_TIMESTAMP_PATTERN

def get_latest_timestamped_file(ext):
    """
    Gets the latest timestamped backup file name with the given database type
    extension.
    """
    pattern = '{dir}/{pattern}.{ext}.gz'.format(
        dir=BACKUP_DIR,
        pattern=BACKUP_TIMESTAMP_PATTERN,
        ext=ext,
    )

    l = glob.glob(pattern)
    l.sort()
    l.reverse()

    if not l:
        raise RestoreError("No backups found matching '{0}' pattern".format(pattern))

    return l[0]
<|endoftext|>


<commit_before>#!/usr/bin/env python
import logging

from mock import patch
import unittest

import boto

from monocyte.loghandler import CloudWatchHandler

class CloudWatchHandlerTest(unittest.TestCase):

    def setUp(self):
        self.boto_mock = patch("monocyte.loghandler.boto").start()
        self.connection = self.boto_mock.logs.connect_to_region.return_value
        self.logger = logging.getLogger("Test")
        self.logger.setLevel(logging.WARN)
        # # do you want a console output, too? use this...
        # console_handler = logging.StreamHandler(sys.stdout)
        # console_handler.setFormatter(logging.Formatter('%(levelname)-8s %(message)s'))
        # self.logger.addHandler(console_handler)

    def myAssertIn(self, a, b):
        self.assertTrue(a in b)

    def test_cloudwatch_logging(self):
        cloudwatch_handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.logger.addHandler(cloudwatch_handler)

        self.logger.debug("aha")
        self.logger.info("aha")
        self.assertFalse(self.connection.put_log_events.called)

        self.logger.warn("aha")

        self.connection.create_log_group.assert_called_with("monocyte")
        self.connection.create_log_stream.assert_called_with("monocyte", "test")
        call_args = self.connection.put_log_events.call_args[0]
        self.assertEqual("monocyte", call_args[0])
        self.assertEqual("test", call_args[1])
        log_event = call_args[2][0]
        self.myAssertIn("timestamp", log_event)
        self.myAssertIn("message", log_event)
        self.myAssertIn("WARNING", log_event["message"])
        self.myAssertIn("aha", log_event["message"])

    def test_log_group_already_exists(self):
        e = boto.logs.exceptions.ResourceAlreadyExistsException(400, "binschonda")
        self.connection.create_log_group.side_effect = e
        self.connection.create_log_stream.side_effect = e
        CloudWatchHandler("eu-central-1", "monocyte", "test")

    def test_response_with_invalid_sequence_token(self):
        e = boto.logs.exceptions.InvalidSequenceTokenException(400, "foo")
        e.body = {"expectedSequenceToken": "correct_token"}
        self.connection.put_log_events.side_effect = e

        handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.assertRaises(boto.logs.exceptions.InvalidSequenceTokenException,
                          handler.put_message, "message", "timestamp")

    def test_invalid_response(self):
        e = boto.logs.exceptions.InvalidSequenceTokenException(500, "completelybroken")
        self.connection.put_log_events.side_effect = e

        handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.assertRaises(boto.logs.exceptions.InvalidSequenceTokenException,
                          handler.put_message, "message", "timestamp")

    def test_response_with_missing_sequence_token(self):
        e = boto.logs.exceptions.InvalidSequenceTokenException(400, "foo")
        e.body = {}
        self.connection.put_log_events.side_effect = e

        handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.assertRaises(boto.logs.exceptions.InvalidSequenceTokenException,
                          handler.put_message, "message", "timestamp")

if __name__ == "__main__":
    unittest.main()
<commit_msg>fixes reference issue.<commit_after>#!/usr/bin/env python
import logging
import boto
import boto.logs.exceptions
import unittest
from mock import patch
from monocyte.loghandler import CloudWatchHandler

class CloudWatchHandlerTest(unittest.TestCase):

    def setUp(self):
        self.boto_mock = patch("monocyte.loghandler.boto").start()
        self.connection = self.boto_mock.logs.connect_to_region.return_value
        self.logger = logging.getLogger("Test")
        self.logger.setLevel(logging.WARN)
        # # do you want a console output, too? use this...
        # console_handler = logging.StreamHandler(sys.stdout)
        # console_handler.setFormatter(logging.Formatter('%(levelname)-8s %(message)s'))
        # self.logger.addHandler(console_handler)

    def myAssertIn(self, a, b):
        self.assertTrue(a in b)

    def test_cloudwatch_logging(self):
        cloudwatch_handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.logger.addHandler(cloudwatch_handler)

        self.logger.debug("aha")
        self.logger.info("aha")
        self.assertFalse(self.connection.put_log_events.called)

        self.logger.warn("aha")

        self.connection.create_log_group.assert_called_with("monocyte")
        self.connection.create_log_stream.assert_called_with("monocyte", "test")
        call_args = self.connection.put_log_events.call_args[0]
        self.assertEqual("monocyte", call_args[0])
        self.assertEqual("test", call_args[1])
        log_event = call_args[2][0]
        self.myAssertIn("timestamp", log_event)
        self.myAssertIn("message", log_event)
        self.myAssertIn("WARNING", log_event["message"])
        self.myAssertIn("aha", log_event["message"])

    def test_log_group_already_exists(self):
        e = boto.logs.exceptions.ResourceAlreadyExistsException(400, "binschonda")
        self.connection.create_log_group.side_effect = e
        self.connection.create_log_stream.side_effect = e
        CloudWatchHandler("eu-central-1", "monocyte", "test")

    def test_response_with_invalid_sequence_token(self):
        e = boto.logs.exceptions.InvalidSequenceTokenException(400, "foo")
        e.body = {"expectedSequenceToken": "correct_token"}
        self.connection.put_log_events.side_effect = e

        handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.assertRaises(boto.logs.exceptions.InvalidSequenceTokenException,
                          handler.put_message, "message", "timestamp")

    def test_invalid_response(self):
        e = boto.logs.exceptions.InvalidSequenceTokenException(500, "completelybroken")
        self.connection.put_log_events.side_effect = e

        handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.assertRaises(boto.logs.exceptions.InvalidSequenceTokenException,
                          handler.put_message, "message", "timestamp")

    def test_response_with_missing_sequence_token(self):
        e = boto.logs.exceptions.InvalidSequenceTokenException(400, "foo")
        e.body = {}
        self.connection.put_log_events.side_effect = e

        handler = CloudWatchHandler("eu-central-1", "monocyte", "test")
        self.assertRaises(boto.logs.exceptions.InvalidSequenceTokenException,
                          handler.put_message, "message", "timestamp")

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before>import datetime
import subprocess
from time import sleep

from pocket_api import Pocket, RateLimitException
from workflow import Workflow, PasswordNotFound
from workflow.background import run_in_background, is_running

from pocket_errors import ERROR_MESSAGES
import config

WF = Workflow(update_settings={
    'github_slug': 'fniephaus/alfred-pocket',
    'version': 'v3.1',
})

def main():
    register_magic_arguments()
    user_input = WF.args[0]

    if WF.update_available:
        WF.add_item(
            "An update is available!",
            autocomplete='workflow:update',
            valid=False
        )

    try:
        WF.get_password('pocket_access_token')
    except PasswordNotFound:
        authorize()

    try:
        WF.get_password('pocket_access_token')
        links = get_links()

        if links:
            if links not in ERROR_MESSAGES.keys():
                add_items(links, user_input)
            else:
                msg = ERROR_MESSAGES[item_list]
                WF.add_item(msg[0], msg[1], valid=False)
        else:
            WF.add_item('Your Pocket list is empty!', valid=False)

        # Update Pocket list in background
        if not WF.cached_data_fresh('pocket_list', max_age=10):
            refresh_list()

    except PasswordNotFound:
        subprocess.call(['open', get_auth_url()])

    WF.send_feedback()

def register_magic_arguments():
    WF.magic_prefix = 'wf:'
    def delete_access_token():
        WF.delete_password('pocket_access_token')
        return 'Access token has been deleted successfully.'
    WF.magic_arguments['deauth'] = delete_access_token

def get_links():
    links = WF.cached_data('pocket_list', max_age=120)
    tries = 10
    # Wait for data
    while(links == None):
        refresh_list()
        sleep(0.5)
        links = WF.cached_data('pocket_list', max_age=120)
        if tries > 0:
            tries -= 1
        else:
            return None
    return links

def add_items(links, user_input):
    links = sorted(links.values(), key=lambda x: int(x['item_id']), reverse=True)
    for index, link in enumerate(links):
        required_keys = [
            'item_id', 'given_title', 'resolved_url', 'time_added']
        if all(x in link for x in required_keys):
            # prepare title
            if len(link.get('resolved_title', '')) > 0:
                title = link['resolved_title']
            else:
                title = link['given_title']
            # prepare subtitle
            link_count = len(links) - index
            tags = link['tags'] if 'tags' in link else None
            subtitle = get_subtitle(
                link_count,
                link['time_added'],
                link['resolved_url'],
                tags
            )
            # prepare argument
            argument = '%s %s' % (
                link['resolved_url'], link['item_id'])

            if (user_input.lower() in title.lower() or
                    user_input.lower() in subtitle.lower()):
                WF.add_item(
                    title,
                    subtitle,
                    arg=argument,
                    valid=True
                )

    if WF._items == []:
        WF.add_item(
            'No links found for "%s".' % user_input,
            valid=False
        )

def get_subtitle(item_count, time_added, resolved_url, tags=None):
    time_updated = datetime.datetime.fromtimestamp(
        int(time_added)).strftime('%Y-%m-%d %H:%M')
    short_url = resolved_url.replace(
        'http://', '').replace('https://', '')

    subtitle_elements = ['#%s' % item_count, time_updated, short_url]

    if tags:
        tags = ['#%s' % x for x in tags.keys()]
        subtitle_elements.insert(2, ', '.join(tags))

    return ' - '.join(subtitle_elements)

def get_auth_url():
    request_token = Pocket.get_request_token(
        consumer_key=config.CONSUMER_KEY, redirect_uri=config.REDIRECT_URI)
    WF.cache_data('pocket_request_token', request_token)

    auth_url = Pocket.get_auth_url(
        code=request_token, redirect_uri=config.REDIRECT_URI)

    return auth_url

def authorize():
    request_token = WF.cached_data('pocket_request_token')
    if request_token:
        try:
            user_credentials = Pocket.get_credentials(
                consumer_key=config.CONSUMER_KEY,
                code=request_token
            )
            WF.save_password(
                'pocket_access_token',
                user_credentials['access_token']
            )
            # We don't need the cache anymore. Clear it for security reasons
            WF.clear_cache()
        except RateLimitException:
            WF.logger.error('RateLimitException')

def refresh_list():
    if not is_running('pocket_refresh'):
        cmd = ['/usr/bin/python', WF.workflowfile('pocket_refresh.py')]
        run_in_background('pocket_refresh', cmd)

if __name__ == '__main__':
    main()
<commit_msg>Sort by time_updated<commit_after>import datetime
import subprocess
from time import sleep

from pocket_api import Pocket, RateLimitException
from workflow import Workflow, PasswordNotFound
from workflow.background import run_in_background, is_running

from pocket_errors import ERROR_MESSAGES
import config

WF = Workflow(update_settings={
    'github_slug': 'fniephaus/alfred-pocket',
    'version': 'v3.1',
})

def main():
    register_magic_arguments()
    user_input = WF.args[0]

    if WF.update_available:
        WF.add_item(
            "An update is available!",
            autocomplete='workflow:update',
            valid=False
        )

    try:
        WF.get_password('pocket_access_token')
    except PasswordNotFound:
        authorize()

    try:
        WF.get_password('pocket_access_token')
        links = get_links()

        if links:
            if links not in ERROR_MESSAGES.keys():
                add_items(links, user_input)
            else:
                msg = ERROR_MESSAGES[item_list]
                WF.add_item(msg[0], msg[1], valid=False)
        else:
            WF.add_item('Your Pocket list is empty!', valid=False)

        # Update Pocket list in background
        if not WF.cached_data_fresh('pocket_list', max_age=10):
            refresh_list()

    except PasswordNotFound:
        subprocess.call(['open', get_auth_url()])

    WF.send_feedback()

def register_magic_arguments():
    WF.magic_prefix = 'wf:'
    def delete_access_token():
        WF.delete_password('pocket_access_token')
        return 'Access token has been deleted successfully.'
    WF.magic_arguments['deauth'] = delete_access_token

def get_links():
    links = WF.cached_data('pocket_list', max_age=120)
    tries = 10
    # Wait for data
    while(links == None):
        refresh_list()
        sleep(0.5)
        links = WF.cached_data('pocket_list', max_age=120)
        if tries > 0:
            tries -= 1
        else:
            return None
    return links

def add_items(links, user_input):
    links = sorted(links.values(), key=lambda x: int(x['time_updated']), reverse=True)
    for index, link in enumerate(links):
        required_keys = [
            'item_id', 'given_title', 'resolved_url', 'time_updated']
        if all(x in link for x in required_keys):
            # prepare title
            if len(link.get('resolved_title', '')) > 0:
                title = link['resolved_title']
            else:
                title = link['given_title']
            # prepare subtitle
            link_count = len(links) - index
            tags = link['tags'] if 'tags' in link else None
            subtitle = get_subtitle(
                link_count,
                link['time_updated'],
                link['resolved_url'],
                tags
            )
            # prepare argument
            argument = '%s %s' % (
                link['resolved_url'], link['item_id'])

            if (user_input.lower() in title.lower() or
                    user_input.lower() in subtitle.lower()):
                WF.add_item(
                    title,
                    subtitle,
                    arg=argument,
                    valid=True
                )

    if WF._items == []:
        WF.add_item(
            'No links found for "%s".' % user_input,
            valid=False
        )

def get_subtitle(item_count, time_updated, resolved_url, tags=None):
    time_updated = datetime.datetime.fromtimestamp(
        int(time_updated)).strftime('%Y-%m-%d %H:%M')
    short_url = resolved_url.replace(
        'http://', '').replace('https://', '')

    subtitle_elements = ['#%s' % item_count, time_updated, short_url]

    if tags:
        tags = ['#%s' % x for x in tags.keys()]
        subtitle_elements.insert(2, ', '.join(tags))

    return ' - '.join(subtitle_elements)

def get_auth_url():
    request_token = Pocket.get_request_token(
        consumer_key=config.CONSUMER_KEY, redirect_uri=config.REDIRECT_URI)
    WF.cache_data('pocket_request_token', request_token)

    auth_url = Pocket.get_auth_url(
        code=request_token, redirect_uri=config.REDIRECT_URI)

    return auth_url

def authorize():
    request_token = WF.cached_data('pocket_request_token')
    if request_token:
        try:
            user_credentials = Pocket.get_credentials(
                consumer_key=config.CONSUMER_KEY,
                code=request_token
            )
            WF.save_password(
                'pocket_access_token',
                user_credentials['access_token']
            )
            # We don't need the cache anymore. Clear it for security reasons
            WF.clear_cache()
        except RateLimitException:
            WF.logger.error('RateLimitException')

def refresh_list():
    if not is_running('pocket_refresh'):
        cmd = ['/usr/bin/python', WF.workflowfile('pocket_refresh.py')]
        run_in_background('pocket_refresh', cmd)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from StringIO import StringIO
from hashlib import md5

from zope.interface.verify import verifyClass, verifyObject
from zope.interface import implements

from twisted.internet.defer import Deferred, succeed, fail
from twisted.python import log
from twisted.trial.unittest import TestCase

from txaws.s3.client import S3Client
from txaws.s3.model import (MultipartInitiationResponse,
    MultipartCompletionResponse)
from txaws.service import AWSServiceRegion
from txaws.credentials import AWSCredentials

from bafload.interfaces import (IPartsGenerator, ITransmissionCounter,
    IPartHandler, IMultipartUploadsManager)
from bafload.errors import UploadPartError
from bafload.up import (FileIOPartsGenerator, PartsTransferredCounter,
    SingleProcessPartUploader, MultipartUploadsManager, MultipartUpload)
from bafload.test.util import FakeLog, FakeS3Client
from bafload import up as up_module

class InterfacesTestCase(TestCase):

    def test_file_io_parts_generator_ifaces(self):
        verifyClass(IPartsGenerator, FileIOPartsGenerator)
        parts_gen = FileIOPartsGenerator()
        verifyObject(IPartsGenerator, parts_gen)

    def test_parts_transmission_counter_ifaces(self):
        verifyClass(ITransmissionCounter, PartsTransferredCounter)
        counter = PartsTransferredCounter(365)
        verifyObject(ITransmissionCounter, counter)

    def test_single_process_part_uploader_ifaces(self):
        verifyClass(IPartHandler, SingleProcessPartUploader)
        uploader = SingleProcessPartUploader()
        verifyObject(IPartHandler, uploader)

    def test_multipart_upload_ifaces(self):
        verifyClass(IMultipartUploadsManager, MultipartUploadsManager)
        mp_uploader = MultipartUploadsManager()
        verifyObject(IMultipartUploadsManager, mp_uploader)

class TestMultipartUpload(MultipartUpload):

    def upload(self, bucket, object_name, content_type, metadata,
               amz_headers={}):
        self.bucket = bucket
        self.object_name = object_name
        self.metadata = metadata
        self.amz_headers = amz_headers
        self.finished.callback(self)

class FileIOPartsGeneratorTestCase(TestCase):

    def setUp(self):
        super(FileIOPartsGeneratorTestCase, self).setUp()
        self.generator = FileIOPartsGenerator()
        self.generator.part_size = 10

    def test_generate_parts(self):
        fd = StringIO("x" * 5 + "y" * 10 + "z" * 13)
        generated = [ entity for entity in self.generator.generate_parts(fd) ]
        self.assertEqual(generated, [('xxxxxyyyyy', 1),
                                     ('yyyyyzzzzz', 2),
                                     ('zzzzzzzz', 3)])

    def test_default_part_size(self):
        self.assertEqual(FileIOPartsGenerator.part_size, 5 * 1024 * 1024)

    def test_count_parts_for_stringio_type(self):
        fd = StringIO("x" * 53)
        count = self.generator.count_parts(fd)
        self.assertEqual(count, 6)
        fd = StringIO("x" * 30)
        count = self.generator.count_parts(fd)
        self.assertEqual(count, 3)

    def test_count_parts_for_file_type(self):
        path = self.mktemp()
        with open(path, "w") as fd:
            fd.write("x" * 53)
        fd = open(path)
        count = self.generator.count_parts(fd)
        self.assertEqual(count, 6)

    def test_count_parts_for_other_type(self):
        count = self.generator.count_parts([])
        self.assertEqual(count, "?")

class DummyPartsGenerator(object):
    implements(IPartsGenerator)

    def __init__(self):
        self.generated = []

    def generate_parts(self, fd):
        for i in range(10):
            entity = (chr(ord('a') + i) * 10, i+1)
            self.generated.append(entity)
            yield entity

class DummyPartHandler(object):

    def __init__(self):
        self.handled = []

    def handle_part(self, part, part_number):
        self.handled.append((part, part_number))
        etag = md5(part).hexdigest()
        return succeed((etag, part_number))

class ErroringPartHandler(object):

    def __init__(self, error_count=3):
        self.handled = []
        self.error_count = error_count
        self.errors = {}

    def handle_part(self, part, part_number):
        self.errors.setdefault(part_number, [])
        if len(self.errors[part_number]) < self.error_count:
            self.errors[part_number].append(1)
            return fail(ValueError('woops'))
        self.handled.append((part, part_number))
        etag = md5(part).hexdigest()
        return succeed((etag, part_number))

class SingleProcessPartUploaderTestCase(TestCase):

    def test_handle_part(self):
        client = FakeS3Client()
        handler = SingleProcessPartUploader()
        handler.bucket = "mybucket"
        handler.object_name = "mykey"
        handler.upload_id = "theupload"
        handler.client = client
        def check(result):
            self.assertEqual(result, (1, '"0123456789"'))
        d = handler.handle_part("aaaaaaaaaa", 1)
        d.addCallback(check)
        return d

class MultipartUploadTestCase(TestCase):

    def setUp(self):
        super(MultipartUploadTestCase, self).setUp()
        self.log = FakeLog()
        self.delayed_calls = []
        self.patch(up_module, 'reactor', self)

    def callLater(self, delay, f, *a, **kw):
        self.delayed_calls.append((delay, f, a, kw))
        f(*a, **kw)

    def test_str(self):
        upload = MultipartUpload(None, None, None, None, None, Deferred())
        r = MultipartInitiationResponse('', '', '1234567890')
        upload.init_response = r
        self.assertEqual(str(upload), "MultipartUpload upload_id=1234567890")

    def test_upload(self):
        client = FakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = DummyPartHandler()
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        def check(task):
            self.assertIdentical(task, upload)
            expected = [('aaaaaaaaaa', 1),
                        ('bbbbbbbbbb', 2),
                        ('cccccccccc', 3),
                        ('dddddddddd', 4),
                        ('eeeeeeeeee', 5),
                        ('ffffffffff', 6),
                        ('gggggggggg', 7),
                        ('hhhhhhhhhh', 8),
                        ('iiiiiiiiii', 9),
                        ('jjjjjjjjjj', 10)]
            self.assertEqual(parts_generator.generated, expected)
            self.assertEqual(part_handler.handled, expected)
            expected = [
                ('e09c80c42fda55f9d992e59ca6b3307d', 1),
                ('82136b4240d6ce4ea7d03e51469a393b', 2),
                ('e9b3390206d8dfc5ffc9b09284c0bbde', 3),
                ('a9e49c7aefe022f0a8540361cce7575c', 4),
                ('c5ba867d9056b7cecf87f8ce88af90f8', 5),
                ('66952c6203ae23242590c0061675234d', 6),
                ('fdc68ea4cf2763996cf215451b291c63', 7),
                ('70270ca63a3de2d8905a9181a0245e58', 8),
                ('d98bcb28df1de541e95a6722c5e983ea', 9),
                ('c71a8da22bf4053760a604897627474c', 10),]
            self.assertEqual(received, expected)
            self.assertIsInstance(task.init_response,
                MultipartInitiationResponse)
            self.assertIsInstance(task.completion_response,
                MultipartCompletionResponse)
            self.assertEqual(client.calls, [
                ('init_multipart_upload', 'mybucket',
                 'mykey', '', {}, {'acl': 'public-read'})])
        d.addCallback(check)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        return d

    def test_upload_error_recovery(self):
        client = FakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = ErroringPartHandler()
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        def check(task):
            errors = self.flushLoggedErrors()
            self.assertEqual(len(self.delayed_calls), 3)
            self.assertIdentical(task, upload)
            expected = [('aaaaaaaaaa', 1),
                        ('bbbbbbbbbb', 2),
                        ('cccccccccc', 3),
                        ('dddddddddd', 4),
                        ('eeeeeeeeee', 5),
                        ('ffffffffff', 6),
                        ('gggggggggg', 7),
                        ('hhhhhhhhhh', 8),
                        ('iiiiiiiiii', 9),
                        ('jjjjjjjjjj', 10)]
            self.assertEqual(parts_generator.generated, expected)
            self.assertEqual(part_handler.handled, expected)
            expected = [
                ('e09c80c42fda55f9d992e59ca6b3307d', 1),
                ('82136b4240d6ce4ea7d03e51469a393b', 2),
                ('e9b3390206d8dfc5ffc9b09284c0bbde', 3),
                ('a9e49c7aefe022f0a8540361cce7575c', 4),
                ('c5ba867d9056b7cecf87f8ce88af90f8', 5),
                ('66952c6203ae23242590c0061675234d', 6),
                ('fdc68ea4cf2763996cf215451b291c63', 7),
                ('70270ca63a3de2d8905a9181a0245e58', 8),
                ('d98bcb28df1de541e95a6722c5e983ea', 9),
                ('c71a8da22bf4053760a604897627474c', 10),]
            self.assertEqual(received, expected)
            self.assertIsInstance(task.init_response,
                MultipartInitiationResponse)
            self.assertIsInstance(task.completion_response,
                MultipartCompletionResponse)
            self.assertEqual(client.calls, [
                ('init_multipart_upload', 'mybucket',
                 'mykey', '', {}, {'acl': 'public-read'})])
        d.addCallback(check)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        return d

    def test_upload_error_timeout_finally(self):
        client = FakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = ErroringPartHandler(100)
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        def eb(why):
            self.assertEquals(len(self.delayed_calls), 13)
            return why
        d.addErrback(eb)
        return self.assertFailure(d, UploadPartError)

class MultipartUploadsManagerTestCase(TestCase):

    def setUp(self):
        super(MultipartUploadsManagerTestCase, self).setUp()
        self.log = FakeLog()
        self.patch(up_module, 'MultipartUpload', TestMultipartUpload)

    def test_region(self):
        region = AWSServiceRegion()
        manager = MultipartUploadsManager(region=region)
        self.assertIdentical(manager.region, region)

    def test_creds(self):
        creds = AWSCredentials()
        manager = MultipartUploadsManager(creds=creds)
        self.assertIdentical(manager.region.creds, creds)

    def test_upload_creation(self):
        def check(task):
            self.assertIsInstance(task, MultipartUpload)
            self.assertEqual(task.bucket, "mybucket")
            self.assertEqual(task.object_name, "mykey")
            self.assertIdentical(task.fd, fd)
            self.assertEqual(task.metadata, {})
            self.assertIsInstance(task.counter, PartsTransferredCounter)
            self.assertIsInstance(task.client, S3Client)
            self.assertEqual(task.amz_headers, {'acl': 'public-read'})
            self.assert_(self.log.buffer)
            verifyObject(IPartsGenerator, task.parts_generator)
            verifyObject(IPartHandler, task.part_handler)
            for entry in self.log.buffer:
                self.assertEqual(entry[0], 'msg')
        manager = MultipartUploadsManager(log=self.log)
        fd = StringIO("some data")
        d = manager.upload(fd, 'mybucket', 'mykey',
                           amz_headers={'acl':'public-read'})
        d.addCallback(check)
        return d
<commit_msg>unit test for retry on error for multipart complete<commit_after>from StringIO import StringIO
from hashlib import md5

from zope.interface.verify import verifyClass, verifyObject
from zope.interface import implements

from twisted.internet.defer import Deferred, succeed, fail
from twisted.python import log
from twisted.trial.unittest import TestCase

from txaws.s3.client import S3Client
from txaws.s3.model import (MultipartInitiationResponse,
    MultipartCompletionResponse)
from txaws.service import AWSServiceRegion
from txaws.credentials import AWSCredentials

from bafload.interfaces import (IPartsGenerator, ITransmissionCounter,
    IPartHandler, IMultipartUploadsManager)
from bafload.errors import UploadPartError
from bafload.up import (FileIOPartsGenerator, PartsTransferredCounter,
    SingleProcessPartUploader, MultipartUploadsManager, MultipartUpload)
from bafload.test.util import FakeLog, FakeS3Client
from bafload import up as up_module

class InterfacesTestCase(TestCase):

    def test_file_io_parts_generator_ifaces(self):
        verifyClass(IPartsGenerator, FileIOPartsGenerator)
        parts_gen = FileIOPartsGenerator()
        verifyObject(IPartsGenerator, parts_gen)

    def test_parts_transmission_counter_ifaces(self):
        verifyClass(ITransmissionCounter, PartsTransferredCounter)
        counter = PartsTransferredCounter(365)
        verifyObject(ITransmissionCounter, counter)

    def test_single_process_part_uploader_ifaces(self):
        verifyClass(IPartHandler, SingleProcessPartUploader)
        uploader = SingleProcessPartUploader()
        verifyObject(IPartHandler, uploader)

    def test_multipart_upload_ifaces(self):
        verifyClass(IMultipartUploadsManager, MultipartUploadsManager)
        mp_uploader = MultipartUploadsManager()
        verifyObject(IMultipartUploadsManager, mp_uploader)

class TestMultipartUpload(MultipartUpload):

    def upload(self, bucket, object_name, content_type, metadata,
               amz_headers={}):
        self.bucket = bucket
        self.object_name = object_name
        self.metadata = metadata
        self.amz_headers = amz_headers
        self.finished.callback(self)

class FileIOPartsGeneratorTestCase(TestCase):

    def setUp(self):
        super(FileIOPartsGeneratorTestCase, self).setUp()
        self.generator = FileIOPartsGenerator()
        self.generator.part_size = 10

    def test_generate_parts(self):
        fd = StringIO("x" * 5 + "y" * 10 + "z" * 13)
        generated = [ entity for entity in self.generator.generate_parts(fd) ]
        self.assertEqual(generated, [('xxxxxyyyyy', 1),
                                     ('yyyyyzzzzz', 2),
                                     ('zzzzzzzz', 3)])

    def test_default_part_size(self):
        self.assertEqual(FileIOPartsGenerator.part_size, 5 * 1024 * 1024)

    def test_count_parts_for_stringio_type(self):
        fd = StringIO("x" * 53)
        count = self.generator.count_parts(fd)
        self.assertEqual(count, 6)
        fd = StringIO("x" * 30)
        count = self.generator.count_parts(fd)
        self.assertEqual(count, 3)

    def test_count_parts_for_file_type(self):
        path = self.mktemp()
        with open(path, "w") as fd:
            fd.write("x" * 53)
        fd = open(path)
        count = self.generator.count_parts(fd)
        self.assertEqual(count, 6)

    def test_count_parts_for_other_type(self):
        count = self.generator.count_parts([])
        self.assertEqual(count, "?")

class DummyPartsGenerator(object):
    implements(IPartsGenerator)

    def __init__(self):
        self.generated = []

    def generate_parts(self, fd):
        for i in range(10):
            entity = (chr(ord('a') + i) * 10, i+1)
            self.generated.append(entity)
            yield entity

class DummyPartHandler(object):

    def __init__(self):
        self.handled = []

    def handle_part(self, part, part_number):
        self.handled.append((part, part_number))
        etag = md5(part).hexdigest()
        return succeed((etag, part_number))

class ErroringPartHandler(object):

    def __init__(self, error_count=3):
        self.handled = []
        self.error_count = error_count
        self.errors = {}

    def handle_part(self, part, part_number):
        self.errors.setdefault(part_number, [])
        if len(self.errors[part_number]) < self.error_count:
            self.errors[part_number].append(1)
            return fail(ValueError('woops'))
        self.handled.append((part, part_number))
        etag = md5(part).hexdigest()
        return succeed((etag, part_number))

class ErroringFakeS3Client(FakeS3Client):

    def __init__(self, error_count=3):
        FakeS3Client.__init__(self)
        self.error_count = error_count

    def complete_multipart_upload(self, *p, **kw):
        if not self.error_count:
            return FakeS3Client.complete_multipart_upload(self, *p, **kw)
        self.error_count -= 1
        return fail(ValueError('woops'))

class SingleProcessPartUploaderTestCase(TestCase):

    def test_handle_part(self):
        client = FakeS3Client()
        handler = SingleProcessPartUploader()
        handler.bucket = "mybucket"
        handler.object_name = "mykey"
        handler.upload_id = "theupload"
        handler.client = client
        def check(result):
            self.assertEqual(result, (1, '"0123456789"'))
        d = handler.handle_part("aaaaaaaaaa", 1)
        d.addCallback(check)
        return d

class MultipartUploadTestCase(TestCase):

    def setUp(self):
        super(MultipartUploadTestCase, self).setUp()
        self.log = FakeLog()
        self.delayed_calls = []
        self.patch(up_module, 'reactor', self)

    def callLater(self, delay, f, *a, **kw):
        self.delayed_calls.append((delay, f, a, kw))
        f(*a, **kw)

    def test_str(self):
        upload = MultipartUpload(None, None, None, None, None, Deferred())
        r = MultipartInitiationResponse('', '', '1234567890')
        upload.init_response = r
        self.assertEqual(str(upload), "MultipartUpload upload_id=1234567890")

    def _assertPartsCompleted(self, parts_generator, part_handler, received,
                              task, client):
        expected = [('aaaaaaaaaa', 1),
                    ('bbbbbbbbbb', 2),
                    ('cccccccccc', 3),
                    ('dddddddddd', 4),
                    ('eeeeeeeeee', 5),
                    ('ffffffffff', 6),
                    ('gggggggggg', 7),
                    ('hhhhhhhhhh', 8),
                    ('iiiiiiiiii', 9),
                    ('jjjjjjjjjj', 10)]
        self.assertEqual(parts_generator.generated, expected)
        self.assertEqual(part_handler.handled, expected)
        expected = [
            ('e09c80c42fda55f9d992e59ca6b3307d', 1),
            ('82136b4240d6ce4ea7d03e51469a393b', 2),
            ('e9b3390206d8dfc5ffc9b09284c0bbde', 3),
            ('a9e49c7aefe022f0a8540361cce7575c', 4),
            ('c5ba867d9056b7cecf87f8ce88af90f8', 5),
            ('66952c6203ae23242590c0061675234d', 6),
            ('fdc68ea4cf2763996cf215451b291c63', 7),
            ('70270ca63a3de2d8905a9181a0245e58', 8),
            ('d98bcb28df1de541e95a6722c5e983ea', 9),
            ('c71a8da22bf4053760a604897627474c', 10),]
        self.assertEqual(received, expected)
        self.assertIsInstance(task.init_response,
            MultipartInitiationResponse)
        self.assertIsInstance(task.completion_response,
            MultipartCompletionResponse)
        self.assertEqual(client.calls, [
            ('init_multipart_upload', 'mybucket',
             'mykey', '', {}, {'acl': 'public-read'})])

    def test_upload(self):
        client = FakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = DummyPartHandler()
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        def check(task):
            self.assertIdentical(task, upload)
            self._assertPartsCompleted(parts_generator, part_handler,
                                      received, task, client)
        d.addCallback(check)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        return d

    def test_upload_error_recovery(self):
        client = FakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = ErroringPartHandler()
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        def check(task):
            self.flushLoggedErrors()
            self.assertEqual(len(self.delayed_calls), 3)
            self.assertIdentical(task, upload)
            self._assertPartsCompleted(parts_generator, part_handler,
                                      received, task, client)
        d.addCallback(check)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        return d

    def test_upload_error_timeout_finally(self):
        client = FakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = ErroringPartHandler(100)
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        def eb(why):
            self.assertEquals(len(self.delayed_calls), 13)
            return why
        d.addErrback(eb)
        return self.assertFailure(d, UploadPartError)

    def test_upload_on_complete_error_recovery(self):
        client = ErroringFakeS3Client()
        parts_generator = DummyPartsGenerator()
        part_handler = DummyPartHandler()
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        def check(task):
            self.flushLoggedErrors()
            self.assertEqual(len(self.delayed_calls), 3)
            self.assertIdentical(task, upload)
            self._assertPartsCompleted(parts_generator, part_handler,
                                      received, task, client)
        d.addCallback(check)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        return d

    def test_upload_on_complete_error_timeout_finally(self):
        client = ErroringFakeS3Client(100)
        parts_generator = DummyPartsGenerator()
        part_handler = DummyPartHandler()
        counter = PartsTransferredCounter('?')
        d = Deferred()
        received = []
        amz_headers = {'acl': 'public-read'}
        upload = MultipartUpload(client, None, parts_generator, part_handler,
            counter, d, self.log)
        received = []
        upload.on_part_generated = received.append
        upload.upload('mybucket', 'mykey', '', {}, amz_headers)
        def eb(why):
            self.assertEquals(len(self.delayed_calls), 13)
            return why
        d.addErrback(eb)
        return self.assertFailure(d, ValueError)

class MultipartUploadsManagerTestCase(TestCase):

    def setUp(self):
        super(MultipartUploadsManagerTestCase, self).setUp()
        self.log = FakeLog()
        self.patch(up_module, 'MultipartUpload', TestMultipartUpload)

    def test_region(self):
        region = AWSServiceRegion()
        manager = MultipartUploadsManager(region=region)
        self.assertIdentical(manager.region, region)

    def test_creds(self):
        creds = AWSCredentials()
        manager = MultipartUploadsManager(creds=creds)
        self.assertIdentical(manager.region.creds, creds)

    def test_upload_creation(self):
        def check(task):
            self.assertIsInstance(task, MultipartUpload)
            self.assertEqual(task.bucket, "mybucket")
            self.assertEqual(task.object_name, "mykey")
            self.assertIdentical(task.fd, fd)
            self.assertEqual(task.metadata, {})
            self.assertIsInstance(task.counter, PartsTransferredCounter)
            self.assertIsInstance(task.client, S3Client)
            self.assertEqual(task.amz_headers, {'acl': 'public-read'})
            self.assert_(self.log.buffer)
            verifyObject(IPartsGenerator, task.parts_generator)
            verifyObject(IPartHandler, task.part_handler)
            for entry in self.log.buffer:
                self.assertEqual(entry[0], 'msg')
        manager = MultipartUploadsManager(log=self.log)
        fd = StringIO("some data")
        d = manager.upload(fd, 'mybucket', 'mykey',
                           amz_headers={'acl':'public-read'})
        d.addCallback(check)
        return d
<|endoftext|>


<commit_before># coding: utf-8
# Copyright 2013 The Font Bakery Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# See AUTHORS.txt for the list of Authors and LICENSE.txt for the License.

from flask import (Blueprint, request, flash, g, session, redirect,
    url_for, current_app)

from .models import User
from ..extensions import db, github
from flask.ext.babel import gettext as _

gitauth = Blueprint('gitauth', __name__, url_prefix='/auth')

@gitauth.after_request
def after_request(response):
    return response

@gitauth.route('/me')
def me():
    # Visit this URL and log as 1st user in database. Usefull if you need to
    # work offline. To change list of github users allowed to workin in single
    # user mode change GITAUTH_LOGIN_LIST config property
    # You need to login using GitHub once before able to use this functionality
    # offline. To add record to database.
    # Only working if server in debug mode.

    if current_app.debug:
        # pylint:disable-msg=E1101
        user = User.get_or_init('offline')
        if user.id and user.login in current_app.config['GITAUTH_LOGIN_LIST']:
            session['user_id'] = user.id
            g.user = user
            flash(_('Welcome!'))

    return redirect(url_for('frontend.splash'))

@gitauth.route('/login')
def login():
    if session.get('user_id', None) is None:
        redirect_uri = url_for('.authorized',
            next=request.args.get('next') or request.referrer or None,
            _external=True)
        params = {'redirect_uri': redirect_uri, 'scope': 'user:email,public_repo'}
        return redirect(github.get_authorize_url(**params))
    else:
        flash(_('Already logged in'))
        return redirect(url_for('frontend.splash'))

@gitauth.route('/callback')
def authorized(next = None):
    next_url = request.args.get('next') or url_for('frontend.splash')

    if not 'code' in request.args:
        flash(_('You did not authorize the request'))
        return redirect(next_url)

    redirect_uri = url_for('.authorized', _external=True)
    data = dict(code=request.args['code'], redirect_uri=redirect_uri)
    auth = github.get_auth_session(data=data)

    token = auth.access_token
    me = auth.get('user').json()

    user = User.get_or_init(me['login'])

    # if user.id is None:
    #     new record isn't saved yet
    #     flash(_('Welcome to Bakery.'))

    # update user data
    user.name = me.get('name', me.get('login'))
    user.github_access_token = token
    user.avatar = me['gravatar_id']
    user.email = me['email']

    # save to database
    db.session.add(user)
    db.session.commit()

    session['user_id'] = user.id
    g.user = user
    flash(_('Welcome!'))

    return redirect(next_url)

@gitauth.route('/logout')
def logout():
    session.pop('user_id', None)
    return redirect(url_for('frontend.splash'))

<commit_msg>Minor: Improve docstring and flash for gitauth me()<commit_after># coding: utf-8
# Copyright 2013 The Font Bakery Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# See AUTHORS.txt for the list of Authors and LICENSE.txt for the License.

from flask import (Blueprint, request, flash, g, session, redirect,
    url_for, current_app)

from .models import User
from ..extensions import db, github
from flask.ext.babel import gettext as _

gitauth = Blueprint('gitauth', __name__, url_prefix='/auth')

@gitauth.after_request
def after_request(response):
    return response

@gitauth.route('/me')
def me():
    """
    Bypass Github authentication and login as 1st user in the database.

    Usage:
        Visit /auth/me to log in as the 1st user in the database, to
        work offline as that user. To allow a github username to work
        in this single user mode, change GITAUTH_LOGIN_LIST config property.
        You need to login using GitHub at least once before this will work.
        This only works if server is in debug mode.
    """
    if current_app.debug:
        # pylint:disable-msg=E1101
        user = User.get_or_init('offline')
        if user.id and user.login in current_app.config['GITAUTH_LOGIN_LIST']:
            session['user_id'] = user.id
            g.user = user
            flash(_('Welcome to single user mode!'))
    return redirect(url_for('frontend.splash'))

@gitauth.route('/login')
def login():
    if session.get('user_id', None) is None:
        redirect_uri = url_for('.authorized',
            next=request.args.get('next') or request.referrer or None,
            _external=True)
        params = {'redirect_uri': redirect_uri, 'scope': 'user:email,public_repo'}
        return redirect(github.get_authorize_url(**params))
    else:
        flash(_('Already logged in'))
        return redirect(url_for('frontend.splash'))

@gitauth.route('/callback')
def authorized(next = None):
    next_url = request.args.get('next') or url_for('frontend.splash')

    if not 'code' in request.args:
        flash(_('You did not authorize the request'))
        return redirect(next_url)

    redirect_uri = url_for('.authorized', _external=True)
    data = dict(code=request.args['code'], redirect_uri=redirect_uri)
    auth = github.get_auth_session(data=data)

    token = auth.access_token
    me = auth.get('user').json()

    user = User.get_or_init(me['login'])

    # if user.id is None:
    #     new record isn't saved yet
    #     flash(_('Welcome to Bakery.'))

    # update user data
    user.name = me.get('name', me.get('login'))
    user.github_access_token = token
    user.avatar = me['gravatar_id']
    user.email = me['email']

    # save to database
    db.session.add(user)
    db.session.commit()

    session['user_id'] = user.id
    g.user = user
    flash(_('Welcome!'))

    return redirect(next_url)

@gitauth.route('/logout')
def logout():
    session.pop('user_id', None)
    return redirect(url_for('frontend.splash'))

<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

'''
A very forgiving JSON parser.
'''

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import codecs
import re

from six import iteritems
from pyparsing import *

__version__ = '0.1.1'

class IllegalValue(object):
    '''
    A value that is illegal in JSON.

    ``parse`` wraps anything that isn't standard JSON into an
    ``IllegalValue`` instance. By default, these are then automatically
    resolved into standard Python types via ``resolve``. However, if you
    pass ``resolver=None`` to ``parse`` and your input contains illegal
    values then your output will contain instances of this class.

    The part of the source that is represented by this instance is
    stored in the ``source`` attribute. That may be the empty string in
    cases like ``[1, , 2]``.
    '''
    def __init__(self, source):
        '''
        Constructor.

        ``source`` is a string.
        '''
        self.source = source

    def __str__(self):
        return source

    def __repr__(self):
        return '<{} {!r}>'.format(self.__class__.__name__, self.source)

    def __eq__(self, other):
        return (isinstance(other, self.__class__) and
                other.source == self.source)

    def __hash__(self):
        return hash(self.source)

def to_float(t):
    '''
    Convert a string that looks like a float to a float.
    '''
    return float(re.sub(r'\s*', '', t[0]))

def to_int(t):
    '''
    Convert a string that looks like an int to an int.
    '''
    return int(re.sub(r'\s*', '', t[0]))

# From http://stackoverflow.com/a/24519338/857390

ESCAPE_SEQUENCE_RE = re.compile(r'''
    ( \\U........      # 8-digit hex escapes
    | \\u....          # 4-digit hex escapes
    | \\x..            # 2-digit hex escapes
    | \\[0-7]{1,3}     # Octal escapes
    | \\N\{[^}]+\}     # Unicode characters by name
    | \\[\\'"abfnrtv]  # Single-character escapes
    )''',
    re.UNICODE | re.VERBOSE
)

def decode_escapes(s):
    '''
    Decode string escape sequences.
    '''
    def decode_match(match):
        return codecs.decode(match.group(0), 'unicode-escape')
    result = ESCAPE_SEQUENCE_RE.sub(decode_match, s)
    return result

L_BRACKET, R_BRACKET, L_BRACE, R_BRACE, COLON, COMMA = map(Suppress, '[]{}:,')

value = Forward()

null = Keyword('null').setParseAction(lambda: None)

true = Keyword('true').setParseAction(lambda: True)
false = Keyword('false').setParseAction(lambda: False)

int_ = Regex(r'[+-]?\s*\d+').setParseAction(to_int)

FLOAT_RE = r'[+-]?\s*(0(\.\d*)?|([1-9]\d*\.?\d*)|(\.\d+))([Ee][+-]?\d+)?'

float_ = Regex(FLOAT_RE).setParseAction(to_float)

# PyParsing's QuotedString preprocesses some escape sequences incorrectly,
# so we're doing things manually.
string_ = Regex(r'".*?(?<!\\)(\\\\)*"')\
            .setParseAction(lambda t: decode_escapes(t[0][1:-1]))

empty = Empty().setParseAction(lambda: IllegalValue(''))

# Special case the empty list to avoid [IllegalValue('')]
list_ = (
    (L_BRACKET + R_BRACKET).setParseAction(lambda: [[]]) |
    Group(L_BRACKET + Optional(delimitedList(value ^ empty)) + R_BRACKET)
)

illegal = Combine(OneOrMore(quotedString ^ Regex(r'[^,{}[\]]')))\
                .setParseAction(lambda t: IllegalValue(t[0].strip()))
# Like ``illegal`` but doesn't allow colons
illegal_key = Combine(OneOrMore(quotedString ^ Regex(r'[^:,{}[\]]')))\
                .setParseAction(lambda t: IllegalValue(t[0].strip()))

dict_key = string_ | illegal_key
dict_value = Optional(value, default=IllegalValue(''))
dict_item = Group(dict_key + Optional(COLON + dict_value,
                  default=IllegalValue('')))
dict_content = delimitedList(dict_item ^ Empty())\
                .setParseAction(lambda t: dict(t.asList()))
dict_ = L_BRACE + dict_content + R_BRACE

value << (dict_ | list_ | float_ | int_ | string_ |
          null | true | false | illegal)

def default_resolver(value):
    '''
    Default resolver for illegal values.

    Case-insensitively resolves

    - ``'true'``, ``'yes'``, and ``'on'`` to ``True``

    - ``'false'``, ``'no'``, and ``'off'`` to ``False``

    - ``'null'`` and ``'none'`` to ``None``

    - ``'inf'``, ``'-inf'``, and ``'nan'`` to their ``float``
      counterparts

    All other strings are returned unchanged.
    '''
    low = value.lower()
    if low in ['true', 'yes', 'on']:
        return True
    if low in ['false', 'no', 'off']:
        return False
    if low in ['null', 'none']:
        return None
    try:
        # Handles +/- NaN/Inf
        return float(re.sub(r'\s+', '', low))
    except ValueError:
        pass
    return decode_escapes(value)

def resolve(data, resolver=default_resolver):
    '''
    Recursively resolve illegal values.

    ``data`` is a potentially nested Python value as returned by
    ``parse``.

    ``resolver`` is a function that maps strings to arbitrary values.

    All instances of ``IllegalValue`` in ``data`` are replaced by the
    result of feeding them into ``resolver``.
    '''
    if isinstance(data, list):
        items = []
        for item in data:
            items.append(resolve(item, resolver))
        return items
    if isinstance(data, dict):
        items = {}
        for key, value in iteritems(data):
            items[resolve(key, resolver)] = resolve(value, resolver)
        return items
    if isinstance(data, IllegalValue):
        return resolver(data.source)
    return data

def parse(s, resolver=default_resolver):
    '''
    Parse a string that contains barely JSON.

    When values that are illegal in JSON are encountered then they are
    by default heuristicaly resolved to a suitable Python type (see
    ``resolve`` and ``default_resolver``). Set ``resolver`` to a custom
    callback or to a falsy value to modify or disable that mechanism.
    '''
    parsed = value.parseString(s, parseAll=True)
    data = parsed.asList()[0]
    if resolver:
        data = resolve(data, resolver=resolver)
    return data

<commit_msg>Use comprehensions in `resolve`<commit_after>#!/usr/bin/env python
# encoding: utf-8

'''
A very forgiving JSON parser.
'''

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import codecs
import re

from six import iteritems
from pyparsing import *

__version__ = '0.1.1'

class IllegalValue(object):
    '''
    A value that is illegal in JSON.

    ``parse`` wraps anything that isn't standard JSON into an
    ``IllegalValue`` instance. By default, these are then automatically
    resolved into standard Python types via ``resolve``. However, if you
    pass ``resolver=None`` to ``parse`` and your input contains illegal
    values then your output will contain instances of this class.

    The part of the source that is represented by this instance is
    stored in the ``source`` attribute. That may be the empty string in
    cases like ``[1, , 2]``.
    '''
    def __init__(self, source):
        '''
        Constructor.

        ``source`` is a string.
        '''
        self.source = source

    def __str__(self):
        return source

    def __repr__(self):
        return '<{} {!r}>'.format(self.__class__.__name__, self.source)

    def __eq__(self, other):
        return (isinstance(other, self.__class__) and
                other.source == self.source)

    def __hash__(self):
        return hash(self.source)

def to_float(t):
    '''
    Convert a string that looks like a float to a float.
    '''
    return float(re.sub(r'\s*', '', t[0]))

def to_int(t):
    '''
    Convert a string that looks like an int to an int.
    '''
    return int(re.sub(r'\s*', '', t[0]))

# From http://stackoverflow.com/a/24519338/857390

ESCAPE_SEQUENCE_RE = re.compile(r'''
    ( \\U........      # 8-digit hex escapes
    | \\u....          # 4-digit hex escapes
    | \\x..            # 2-digit hex escapes
    | \\[0-7]{1,3}     # Octal escapes
    | \\N\{[^}]+\}     # Unicode characters by name
    | \\[\\'"abfnrtv]  # Single-character escapes
    )''',
    re.UNICODE | re.VERBOSE
)

def decode_escapes(s):
    '''
    Decode string escape sequences.
    '''
    def decode_match(match):
        return codecs.decode(match.group(0), 'unicode-escape')
    result = ESCAPE_SEQUENCE_RE.sub(decode_match, s)
    return result

L_BRACKET, R_BRACKET, L_BRACE, R_BRACE, COLON, COMMA = map(Suppress, '[]{}:,')

value = Forward()

null = Keyword('null').setParseAction(lambda: None)

true = Keyword('true').setParseAction(lambda: True)
false = Keyword('false').setParseAction(lambda: False)

int_ = Regex(r'[+-]?\s*\d+').setParseAction(to_int)

FLOAT_RE = r'[+-]?\s*(0(\.\d*)?|([1-9]\d*\.?\d*)|(\.\d+))([Ee][+-]?\d+)?'

float_ = Regex(FLOAT_RE).setParseAction(to_float)

# PyParsing's QuotedString preprocesses some escape sequences incorrectly,
# so we're doing things manually.
string_ = Regex(r'".*?(?<!\\)(\\\\)*"')\
            .setParseAction(lambda t: decode_escapes(t[0][1:-1]))

empty = Empty().setParseAction(lambda: IllegalValue(''))

# Special case the empty list to avoid [IllegalValue('')]
list_ = (
    (L_BRACKET + R_BRACKET).setParseAction(lambda: [[]]) |
    Group(L_BRACKET + Optional(delimitedList(value ^ empty)) + R_BRACKET)
)

illegal = Combine(OneOrMore(quotedString ^ Regex(r'[^,{}[\]]')))\
                .setParseAction(lambda t: IllegalValue(t[0].strip()))
# Like ``illegal`` but doesn't allow colons
illegal_key = Combine(OneOrMore(quotedString ^ Regex(r'[^:,{}[\]]')))\
                .setParseAction(lambda t: IllegalValue(t[0].strip()))

dict_key = string_ | illegal_key
dict_value = Optional(value, default=IllegalValue(''))
dict_item = Group(dict_key + Optional(COLON + dict_value,
                  default=IllegalValue('')))
dict_content = delimitedList(dict_item ^ Empty())\
                .setParseAction(lambda t: dict(t.asList()))
dict_ = L_BRACE + dict_content + R_BRACE

value << (dict_ | list_ | float_ | int_ | string_ |
          null | true | false | illegal)

def default_resolver(value):
    '''
    Default resolver for illegal values.

    Case-insensitively resolves

    - ``'true'``, ``'yes'``, and ``'on'`` to ``True``

    - ``'false'``, ``'no'``, and ``'off'`` to ``False``

    - ``'null'`` and ``'none'`` to ``None``

    - ``'inf'``, ``'-inf'``, and ``'nan'`` to their ``float``
      counterparts

    All other strings are returned unchanged.
    '''
    low = value.lower()
    if low in ['true', 'yes', 'on']:
        return True
    if low in ['false', 'no', 'off']:
        return False
    if low in ['null', 'none']:
        return None
    try:
        # Handles +/- NaN/Inf
        return float(re.sub(r'\s+', '', low))
    except ValueError:
        pass
    return decode_escapes(value)

def resolve(data, resolver=default_resolver):
    '''
    Recursively resolve illegal values.

    ``data`` is a potentially nested Python value as returned by
    ``parse``.

    ``resolver`` is a function that maps strings to arbitrary values.

    All instances of ``IllegalValue`` in ``data`` are replaced by the
    result of feeding them into ``resolver``.
    '''
    if isinstance(data, list):
        return [resolve(item, resolver) for item in data]
    if isinstance(data, dict):
        return {resolve(key, resolver): resolve(value, resolver) for key, value in iteritems(data)}
    if isinstance(data, IllegalValue):
        return resolver(data.source)
    return data

def parse(s, resolver=default_resolver):
    '''
    Parse a string that contains barely JSON.

    When values that are illegal in JSON are encountered then they are
    by default heuristicaly resolved to a suitable Python type (see
    ``resolve`` and ``default_resolver``). Set ``resolver`` to a custom
    callback or to a falsy value to modify or disable that mechanism.
    '''
    parsed = value.parseString(s, parseAll=True)
    data = parsed.asList()[0]
    if resolver:
        data = resolve(data, resolver=resolver)
    return data

<|endoftext|>