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

import datetime
import argparse
import syslog
import contextlib
import os
import os.path
import sys
import boto3
import json

parser = argparse.ArgumentParser(
    prog="produce-steps-json",
    description="""EMR JSON steps producer for DynamoDB table extraction"""
)

parser.add_argument(
    '-a',
    '--appname',
    type=str,
    default="MYAPP",
    help="Name of the application we are exporting tables for.  Used in the S3 path where the dumps finally end up."
)

parser.add_argument(
    '-r',
    '--region',
    type=str,
    default="us-east-1",
    help="The region to connect to for exporting."
)

parser.add_argument(
    '-e',
    '--readtput',
    type=str,
    default="0.25",
    help="The percentage of read throughput to utilize when exporting (def: 0.25)."
)

parser.add_argument(
    '-w',
    '--writetput',
    type=str,
    default="0.5",
    help="The percentage of write throughput to utilize when importing data (def: 0.5)."
)

parser.add_argument(
    '-s',
    '--spikedread',
    type=str,
    help="The value to spike read throughput to before table export"
)

parser.add_argument(
    '-x',
    '--excludes',
    type=str,
    help="A file containing a list of tables to exclude"
)

parser.add_argument(
    '-f',
    '--filter',
    type=str,
    default="",
    help="Only export tables with this filter criteria in the table name."
)

parser.add_argument(
    '-c',
    '--clusters',
    type=str,
    default="",
    help="Number of Clusters to create. (tables will be split evenly among them)"
)

parser.add_argument(
    'destination',
    type=str,
    help="where to place the EMR export and import steps files"
)

parser.add_argument(
    's3location',
    type=str,
    help="The S3 FOLDER path to place export files in and read import files from."
)

def myLog(message):
    procName = __file__
    currDTS = datetime.datetime.now()
    dateTimeStr = currDTS.strftime('%Y/%B/%d/ %H:%M:%S')

    syslogMsg = procName + ": " + message
    syslog.syslog(syslogMsg)
    print('%s %s' % (dateTimeStr, message))

def main(region, filter, destination, writetput, readtput,
         spikedread, s3location, appname, excludes, clusters):

    dateStr = datetime.datetime.now().strftime("%Y/%m/%d/%H_%M.%S")

    conn = boto3.client('dynamodb', region_name=region)

    # Have we been given an excludes file?  If so, read it.  Any tables in here
    # will not have export steps generated for them
    if excludes:
        myLog("excludes specified - reading " + excludes)

        if os.path.exists(excludes):
            exclude_table_list = [line.rstrip('\n') for line in open(excludes)]
        else:
            myLog("Unable to open " + excludes + " for reading")

    if conn:
        myLog("connected to dynamodb (region: %s)" % region)
        myLog("exporting all tables where table name contains %s " % filter)

        # get a list of all tables in the region
        table_list = listTables(conn)
        # print('Table list:\n%s' % json.dumps(table_list, indent=4))

        filtered_list = [x for x in table_list if filter in x]
        # print("Filtered list:\n" + json.dumps(filtered_list, indent=4))

        # TODO: Print out excluded list?

        filtered_excluded_list = [
            x for x in filtered_list if x not in exclude_table_list]
        # print("Filtered and Excluded list:\n" + json.dumps(filtered_excluded_list, indent=4))

        if clusters > 1:
            myLog('Asked to break list of tables up into %s clusters' % str(clusters))
            table_list_list = chunkIt(filtered_excluded_list, clusters)

        myLog('Broke list of tables up into the following clusters:\n%s' % json.dumps(table_list_list, indent=4))

        cluster_number = 1
        for chunk in table_list_list:
            exportSteps = []
            importSteps = []

            # Process this chunk of tables
            myLog("Creating steps for Cluster %02d" % int(cluster_number))

            table_desc_list = describeTables(conn, chunk)

            # Get the path we will use for 'this' backup
            s3ExportPath = generateS3Path(s3location, region, dateStr, appname)

            # Get the path to the update-throughput script
            s3ScriptPath = s3location.rstrip(
                '/') + "/scripts/update-throughput.sh"

            S3PathFilename = destination + "/s3path.info"
            writeFile(s3ExportPath, S3PathFilename)

            # Now process them, ignoring any that don't match our filter
            for table in table_desc_list:
                myLog(
                    "Generating EMR export JSON for table: [%s]" %
                    table['name'])

                autoscale_min_spike_read_capacity = None  # Assume no autoscaling
                autoscale_min_reset_read_capacity = None
                tableS3Path = s3ExportPath + "/" + table['name']

                # Check if table is set to On Demand Capacity
                if int(table['read']) > 0:
                    myLog(
                        "Table uses provisioned capacity - need to add throughput spike and reset steps")
                    # Does this table have autoscaling enabled?
                    scalable_target_info = scalable_target_exists(
                        region, "table/" + table['name'], "dynamodb:table:ReadCapacityUnits")
                    if scalable_target_info is not None:
                        myLog("Table " + table['name'] +
                              " has autoscaling enabled")
                        autoscale_min_spike_read_capacity = scalable_target_info[0]['MinCapacity'] + int(
                            spikedread)
                        autoscale_min_reset_read_capacity = scalable_target_info[0]['MinCapacity']
                        myLog(
                            "Table " +
                            table['name'] +
                            " has a current AS min capacity of " +
                            str(autoscale_min_reset_read_capacity))

                    if spikedread is not None:
                        tputSpikeStep = generateThroughputUpdateStep(
                            table['name'],
                            "Spike",
                            s3ScriptPath,
                            autoscale_min_spike_read_capacity,
                            autoscale_min_spike_read_capacity,
                            table['write'],
                            region)
                        exportSteps.append(tputSpikeStep)
                else:
                    myLog(
                        "Table uses on-demand capacity - no need for spike and reset throughput steps")

                tableExportStep = generateTableExportStep(
                    table['name'], tableS3Path, readtput)
                exportSteps.append(tableExportStep)

                if int(table['read']) > 0:
                    if spikedread is not None:
                        tputResetStep = generateThroughputUpdateStep(
                            table['name'],
                            "Reset",
                            s3ScriptPath,
                            table['read'],
                            autoscale_min_reset_read_capacity,
                            table['write'],
                            region)
                        exportSteps.append(tputResetStep)

                tableImportStep = generateTableImportStep(
                    table['name'], tableS3Path, writetput)
                importSteps.append(tableImportStep)

            # Now we can write out the import and export steps files
            exportJSON = json.dumps(exportSteps, indent=4)
            exportJSONFilename = "%s/exportSteps_%02d.json" % (
                destination, cluster_number)
            writeFile(exportJSON, exportJSONFilename)

            importJSON = json.dumps(importSteps, indent=4)
            importJSONFilename = "%s/importSteps_%02d.json" % (
                destination, cluster_number)
            writeFile(importJSON, importJSONFilename)
            cluster_number += 1

###########
# Add a JSON entry for a single table throughput update step
###########
def generateThroughputUpdateStep(
        tableName,
        stepName,
        s3Path,
        readtput,
        autoscale_min_throughput,
        writetput,
        region):
    myLog("addThroughputUpdateStep (%s) %s" % (stepName, tableName))

    tputUpdateDict = {}

    if autoscale_min_throughput:
        tputUpdateDict = {
            "Name": stepName +
            " Throughput: " +
            tableName,
            "ActionOnFailure": "CONTINUE",
            "Type": "CUSTOM_JAR",
            "Jar": "s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
            "Args": [
                s3Path,
                region,
                tableName,
                str(readtput),
                str(writetput),
                str(autoscale_min_throughput)]}
    else:
        tputUpdateDict = {
            "Name": stepName +
            " Throughput: " +
            tableName,
            "ActionOnFailure": "CONTINUE",
            "Type": "CUSTOM_JAR",
            "Jar": "s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
            "Args": [
                s3Path,
                region,
                tableName,
                str(readtput),
                str(writetput)]}

    return tputUpdateDict

###########
# Add a JSON entry for a single table export step
###########
def generateTableExportStep(
        tableName, s3Path, readtput, jarPath=None, classPath=None):
    myLog("addTableExportStep %s" % tableName)

    if not jarPath:
        # Default JAR
        jarPath = "s3://dynamodb-emr-us-east-1/emr-ddb-storage-handler/2.1.0/emr-ddb-2.1.0.jar"

    if not classPath:
        # Default ClassPath
        classPath = "org.apache.hadoop.dynamodb.tools.DynamoDbExport"

    tableExportDict = {}
    tableExportDict = {"Name": "Export Table:" + tableName,
                       "ActionOnFailure": "CONTINUE",
                       "Type": "CUSTOM_JAR",
                       "Jar": jarPath,
                       "Args": [classPath,
                                s3Path,
                                tableName,
                                readtput,
                                ]
                       }

    return tableExportDict

###########
# Add a JSON entry for a single table import step
###########
def generateTableImportStep(
        tableName, s3Path, writetput, jarPath=None, classPath=None):
    myLog("addTableImportStep %s" % tableName)

    if not jarPath:
        # Default JAR
        jarPath = "s3://dynamodb-emr-us-east-1/emr-ddb-storage-handler/2.1.0/emr-ddb-2.1.0.jar"

    if not classPath:
        # Default ClassPath
        classPath = "org.apache.hadoop.dynamodb.tools.DynamoDbImport"

    tableImportDict = {"Name": "Import Table:" + tableName,
                       "ActionOnFailure": "CONTINUE",
                       "Type": "CUSTOM_JAR",
                       "Jar": jarPath,
                       "Args": [classPath,
                                s3Path,
                                tableName,
                                writetput
                                ]
                       }

    return tableImportDict

###########
# Generate a formatted S3 path which is used in the export and import steps file
###########
def generateS3Path(basePath, region, dateStr, appname):
    myLog("generateS3Path BASE:%s" % basePath)
    basePath = basePath.rstrip('/')

    s3Path = basePath + "/" + region + "/" + appname + "/" + dateStr
    myLog("S3 path generated is %s" % s3Path)

    return s3Path

def describeTables(conn, table_list):
    table_list_return = []

    for table in table_list:
        table_desc = conn.describe_table(TableName=table)
        table_return = dict()
        table_return['name'] = table
        table_return['read'] = str(
            table_desc['Table']['ProvisionedThroughput']['ReadCapacityUnits'])
        table_return['write'] = str(
            table_desc['Table']['ProvisionedThroughput']['WriteCapacityUnits'])
        table_list_return.append(table_return)

    return table_list_return

###########
# Obtain a list of dynamoDB tables from the current region
###########
def listTables(conn):

    table_list_return = []

    # Get the inital list of tables. boto only returns the first 100 tho....
    table_list = conn.list_tables()

    moreTables = True
    while moreTables:
        if 'LastEvaluatedTableName' in table_list:
            LastEvaluatedTableName = table_list['LastEvaluatedTableName']
            moreTables = True
        else:
            LastEvaluatedTableName = ''
            moreTables = False

        for table_name in table_list['TableNames']:
            table_list_return.append(table_name)

        if LastEvaluatedTableName != '':
            table_list = conn.list_tables(
                ExclusiveStartTableName=LastEvaluatedTableName, Limit=100)

    myLog("Read %d tables from dynamodb" % len(table_list_return))

    return table_list_return

# Checks if a dynamo table has a scalable target (ie. is autoscale enabled?)
def scalable_target_exists(region, resource_id, scalable_dimension):
    response = None
    retval = None

    myLog("Checking if scalable target exists for " +
          resource_id + " for dimension " + scalable_dimension)
    client = boto3.client('application-autoscaling', region_name=region)

    try:
        response = client.describe_scalable_targets(
            ServiceNamespace='dynamodb',
            ResourceIds=[
                resource_id,
            ],
            ScalableDimension=scalable_dimension
        )
    except Exception as e:
        myLog("Failed to describe scalable targets " + str(e))

    if response:
        if response['ScalableTargets']:
            retval = response['ScalableTargets']

    return retval

def chunkIt(seq, num):
    avg = len(seq) / float(num)
    out = []
    last = 0.0

    while last < len(seq):
        out.append(seq[int(last):int(last + avg)])
        last += avg

    return out

def writeFile(content, filename):
    myLog("writeFile %s" % filename)

    text_file = open(filename, "w")
    text_file.write(content)
    text_file.close()

@contextlib.contextmanager
def preserve_cwd():
    cwd = os.getcwd()
    try:
        yield
    finally:
        os.chdir(cwd)

if __name__ == '__main__':
    kwargs = dict(parser.parse_args(sys.argv[1:])._get_kwargs())
    main(**kwargs)
<commit_msg>Minor code cleanup<commit_after>#!/usr/bin/env python

import datetime
import argparse
import syslog
import contextlib
import os
import os.path
import sys
import boto3
import json

parser = argparse.ArgumentParser(
    prog="produce-steps-json",
    description="""EMR JSON steps producer for DynamoDB table extraction"""
)

parser.add_argument(
    '-a',
    '--appname',
    type=str,
    default="MYAPP",
    help="Name of the application we are exporting tables for.  Used in the S3 path where the dumps finally end up."
)

parser.add_argument(
    '-r',
    '--region',
    type=str,
    default="us-east-1",
    help="The region to connect to for exporting."
)

parser.add_argument(
    '-e',
    '--readtput',
    type=str,
    default="0.25",
    help="The percentage of read throughput to utilize when exporting (def: 0.25)."
)

parser.add_argument(
    '-w',
    '--writetput',
    type=str,
    default="0.5",
    help="The percentage of write throughput to utilize when importing data (def: 0.5)."
)

parser.add_argument(
    '-s',
    '--spikedread',
    type=str,
    help="The value to spike read throughput to before table export"
)

parser.add_argument(
    '-x',
    '--excludes',
    type=str,
    help="A file containing a list of tables to exclude"
)

parser.add_argument(
    '-f',
    '--filter',
    type=str,
    default="",
    help="Only export tables with this filter criteria in the table name."
)

parser.add_argument(
    '-c',
    '--clusters',
    type=str,
    default="",
    help="Number of Clusters to create. (tables will be split evenly among them)"
)

parser.add_argument(
    'destination',
    type=str,
    help="where to place the EMR export and import steps files"
)

parser.add_argument(
    's3location',
    type=str,
    help="The S3 FOLDER path to place export files in and read import files from."
)

def myLog(message):
    procName = __file__
    currDTS = datetime.datetime.now()
    dateTimeStr = currDTS.strftime('%Y/%B/%d/ %H:%M:%S')

    syslogMsg = procName + ": " + message
    syslog.syslog(syslogMsg)
    print('%s %s' % (dateTimeStr, message))

def main(region, filter, destination, writetput, readtput,
         spikedread, s3location, appname, excludes, clusters):

    dateStr = datetime.datetime.now().strftime("%Y/%m/%d/%H_%M.%S")

    conn = boto3.client('dynamodb', region_name=region)

    # Have we been given an excludes file?  If so, read it.  Any tables in here
    # will not have export steps generated for them
    if excludes:
        myLog("excludes specified - reading " + excludes)

        if os.path.exists(excludes):
            exclude_table_list = [line.rstrip('\n') for line in open(excludes)]
        else:
            myLog("Unable to open " + excludes + " for reading")

    if conn:
        myLog("Connected to dynamodb (region: %s)" % region)

        # Get the path we will use for 'this' backup
        s3ExportPath = generateS3Path(s3location, region, dateStr, appname)

        # Get the path to the update-throughput script
        s3ScriptPath = s3location.rstrip('/') + "/scripts/update-throughput.sh"

        S3PathFilename = destination + "/s3path.info"
        writeFile(s3ExportPath, S3PathFilename)

        # get a list of all tables in the region
        table_list = listTables(conn)
        # print('Table list:\n%s' % json.dumps(table_list, indent=4))

        myLog("Exporting all tables where table name contains %s " % filter)

        filtered_list = [x for x in table_list if filter in x]
        # print("Filtered list:\n" + json.dumps(filtered_list, indent=4))

        myLog('Excluding any tables in the following list:\n%s' % json.dumps(exclude_table_list, indent=4))

        filtered_excluded_list = [x for x in filtered_list if x not in exclude_table_list]
        # print("Filtered and Excluded list:\n" + json.dumps(filtered_excluded_list, indent=4))

        myLog('Asked to break list of tables up into %s clusters' % str(clusters))
        table_list_list = chunkIt(filtered_excluded_list, clusters)

        myLog('Broke list of tables up into the following clusters:\n%s' % json.dumps(table_list_list, indent=4))

        cluster_number = 1
        for chunk in table_list_list:
            exportSteps = []
            importSteps = []

            # Process this chunk of tables
            myLog("Creating steps for Cluster %02d" % int(cluster_number))

            table_desc_list = describeTables(conn, chunk)

            for table in table_desc_list:
                myLog(
                    "Generating EMR export JSON for table: [%s]" %
                    table['name'])

                autoscale_min_spike_read_capacity = None  # Assume no autoscaling
                autoscale_min_reset_read_capacity = None
                tableS3Path = s3ExportPath + "/" + table['name']

                # Check if table is set to On Demand Capacity
                if int(table['read']) > 0:
                    myLog(
                        "Table uses provisioned capacity - need to add throughput spike and reset steps")
                    # Does this table have autoscaling enabled?
                    scalable_target_info = scalable_target_exists(
                        region, "table/" + table['name'], "dynamodb:table:ReadCapacityUnits")
                    if scalable_target_info is not None:
                        myLog("Table " + table['name'] +
                              " has autoscaling enabled")
                        autoscale_min_spike_read_capacity = scalable_target_info[0]['MinCapacity'] + int(
                            spikedread)
                        autoscale_min_reset_read_capacity = scalable_target_info[0]['MinCapacity']
                        myLog(
                            "Table " +
                            table['name'] +
                            " has a current AS min capacity of " +
                            str(autoscale_min_reset_read_capacity))

                    if spikedread is not None:
                        tputSpikeStep = generateThroughputUpdateStep(
                            table['name'],
                            "Spike",
                            s3ScriptPath,
                            autoscale_min_spike_read_capacity,
                            autoscale_min_spike_read_capacity,
                            table['write'],
                            region)
                        exportSteps.append(tputSpikeStep)
                else:
                    myLog(
                        "Table uses on-demand capacity - no need for spike and reset throughput steps")

                tableExportStep = generateTableExportStep(
                    table['name'], tableS3Path, readtput)
                exportSteps.append(tableExportStep)

                if int(table['read']) > 0:
                    if spikedread is not None:
                        tputResetStep = generateThroughputUpdateStep(
                            table['name'],
                            "Reset",
                            s3ScriptPath,
                            table['read'],
                            autoscale_min_reset_read_capacity,
                            table['write'],
                            region)
                        exportSteps.append(tputResetStep)

                tableImportStep = generateTableImportStep(
                    table['name'], tableS3Path, writetput)
                importSteps.append(tableImportStep)

            # Now we can write out the import and export steps files
            exportJSON = json.dumps(exportSteps, indent=4)
            exportJSONFilename = "%s/exportSteps_%02d.json" % (destination, cluster_number)
            writeFile(exportJSON, exportJSONFilename)

            importJSON = json.dumps(importSteps, indent=4)
            importJSONFilename = "%s/importSteps_%02d.json" % (destination, cluster_number)
            writeFile(importJSON, importJSONFilename)
            cluster_number += 1

###########
# Add a JSON entry for a single table throughput update step
###########
def generateThroughputUpdateStep(
        tableName,
        stepName,
        s3Path,
        readtput,
        autoscale_min_throughput,
        writetput,
        region):
    myLog("addThroughputUpdateStep (%s) %s" % (stepName, tableName))

    tputUpdateDict = {}

    if autoscale_min_throughput:
        tputUpdateDict = {
            "Name": stepName +
            " Throughput: " +
            tableName,
            "ActionOnFailure": "CONTINUE",
            "Type": "CUSTOM_JAR",
            "Jar": "s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
            "Args": [
                s3Path,
                region,
                tableName,
                str(readtput),
                str(writetput),
                str(autoscale_min_throughput)]}
    else:
        tputUpdateDict = {
            "Name": stepName +
            " Throughput: " +
            tableName,
            "ActionOnFailure": "CONTINUE",
            "Type": "CUSTOM_JAR",
            "Jar": "s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
            "Args": [
                s3Path,
                region,
                tableName,
                str(readtput),
                str(writetput)]}

    return tputUpdateDict

###########
# Add a JSON entry for a single table export step
###########
def generateTableExportStep(
        tableName, s3Path, readtput, jarPath=None, classPath=None):
    myLog("addTableExportStep %s" % tableName)

    if not jarPath:
        # Default JAR
        jarPath = "s3://dynamodb-emr-us-east-1/emr-ddb-storage-handler/2.1.0/emr-ddb-2.1.0.jar"

    if not classPath:
        # Default ClassPath
        classPath = "org.apache.hadoop.dynamodb.tools.DynamoDbExport"

    tableExportDict = {"Name": "Export Table:" + tableName,
                       "ActionOnFailure": "CONTINUE",
                       "Type": "CUSTOM_JAR",
                       "Jar": jarPath,
                       "Args": [classPath,
                                s3Path,
                                tableName,
                                readtput,
                                ]
                       }

    return tableExportDict

###########
# Add a JSON entry for a single table import step
###########
def generateTableImportStep(
        tableName, s3Path, writetput, jarPath=None, classPath=None):
    myLog("addTableImportStep %s" % tableName)

    if not jarPath:
        # Default JAR
        jarPath = "s3://dynamodb-emr-us-east-1/emr-ddb-storage-handler/2.1.0/emr-ddb-2.1.0.jar"

    if not classPath:
        # Default ClassPath
        classPath = "org.apache.hadoop.dynamodb.tools.DynamoDbImport"

    tableImportDict = {"Name": "Import Table:" + tableName,
                       "ActionOnFailure": "CONTINUE",
                       "Type": "CUSTOM_JAR",
                       "Jar": jarPath,
                       "Args": [classPath,
                                s3Path,
                                tableName,
                                writetput
                                ]
                       }

    return tableImportDict

###########
# Generate a formatted S3 path which is used in the export and import steps file
###########
def generateS3Path(basePath, region, dateStr, appname):
    myLog("generateS3Path BASE:%s" % basePath)
    basePath = basePath.rstrip('/')

    s3Path = basePath + "/" + region + "/" + appname + "/" + dateStr
    myLog("S3 path generated is %s" % s3Path)

    return s3Path

def describeTables(conn, table_list):
    table_list_return = []

    for table in table_list:
        table_desc = conn.describe_table(TableName=table)
        table_return = dict()
        table_return['name'] = table
        table_return['read'] = str(
            table_desc['Table']['ProvisionedThroughput']['ReadCapacityUnits'])
        table_return['write'] = str(
            table_desc['Table']['ProvisionedThroughput']['WriteCapacityUnits'])
        table_list_return.append(table_return)

    return table_list_return

###########
# Obtain a list of dynamoDB tables from the current region
###########
def listTables(conn):

    table_list_return = []

    # Get the inital list of tables. boto only returns the first 100 tho....
    table_list = conn.list_tables()

    moreTables = True
    while moreTables:
        if 'LastEvaluatedTableName' in table_list:
            LastEvaluatedTableName = table_list['LastEvaluatedTableName']
            moreTables = True
        else:
            LastEvaluatedTableName = ''
            moreTables = False

        for table_name in table_list['TableNames']:
            table_list_return.append(table_name)

        if LastEvaluatedTableName != '':
            table_list = conn.list_tables(
                ExclusiveStartTableName=LastEvaluatedTableName, Limit=100)

    myLog("Read %d tables from dynamodb" % len(table_list_return))

    return table_list_return

# Checks if a dynamo table has a scalable target (ie. is autoscale enabled?)
def scalable_target_exists(region, resource_id, scalable_dimension):
    response = None
    retval = None

    myLog("Checking if scalable target exists for " +
          resource_id + " for dimension " + scalable_dimension)
    client = boto3.client('application-autoscaling', region_name=region)

    try:
        response = client.describe_scalable_targets(
            ServiceNamespace='dynamodb',
            ResourceIds=[
                resource_id,
            ],
            ScalableDimension=scalable_dimension
        )
    except Exception as e:
        myLog("Failed to describe scalable targets " + str(e))

    if response:
        if response['ScalableTargets']:
            retval = response['ScalableTargets']

    return retval

def chunkIt(seq, num):
    avg = len(seq) / float(num)
    out = []
    last = 0.0

    while last < len(seq):
        out.append(seq[int(last):int(last + avg)])
        last += avg

    return out

def writeFile(content, filename):
    myLog("writeFile %s" % filename)

    text_file = open(filename, "w")
    text_file.write(content)
    text_file.close()

@contextlib.contextmanager
def preserve_cwd():
    cwd = os.getcwd()
    try:
        yield
    finally:
        os.chdir(cwd)

if __name__ == '__main__':
    kwargs = dict(parser.parse_args(sys.argv[1:])._get_kwargs())
    main(**kwargs)
<|endoftext|>


<commit_before>import asyncio
import logging
import threading
import time
import uvloop
import unittest
import weakref

from unittest import mock
from uvloop._testbase import UVTestCase, AIOTestCase

class _TestBase:

    def test_close(self):
        self.assertFalse(self.loop._closed)
        self.assertFalse(self.loop.is_closed())
        self.loop.close()
        self.assertTrue(self.loop._closed)
        self.assertTrue(self.loop.is_closed())

        # it should be possible to call close() more than once
        self.loop.close()
        self.loop.close()

        # operation blocked when the loop is closed
        f = asyncio.Future(loop=self.loop)
        self.assertRaises(RuntimeError, self.loop.run_forever)
        self.assertRaises(RuntimeError, self.loop.run_until_complete, f)

    def test_handle_weakref(self):
        wd = weakref.WeakValueDictionary()
        h = self.loop.call_soon(lambda: None)
        wd['h'] = h  # Would fail without __weakref__ slot.

    def test_call_soon(self):
        calls = []

        def cb(inc):
            calls.append(inc)
            self.loop.stop()

        self.loop.call_soon(cb, 10)

        h = self.loop.call_soon(cb, 100)
        self.assertIn('.cb', repr(h))
        h.cancel()
        self.assertIn('cancelled', repr(h))

        self.loop.call_soon(cb, 1)

        self.loop.run_forever()

        self.assertEqual(calls, [10, 1])

    def test_call_soon_base_exc(self):
        def cb():
            raise KeyboardInterrupt()

        self.loop.call_soon(cb)

        with self.assertRaises(KeyboardInterrupt):
            self.loop.run_forever()

        self.assertFalse(self.loop.is_closed())

    def test_calls_debug_reporting(self):
        def run_test(debug, meth, stack_adj):
            context = None

            def handler(loop, ctx):
                nonlocal context
                context = ctx

            self.loop.set_debug(debug)
            self.loop.set_exception_handler(handler)

            def cb():
                1 / 0

            meth(cb)
            self.assertIsNone(context)
            self.loop.run_until_complete(asyncio.sleep(0.05, loop=self.loop))

            self.assertIs(type(context['exception']), ZeroDivisionError)
            self.assertTrue(context['message'].startswith(
                'Exception in callback'))

            if debug:
                tb = context['source_traceback']
                self.assertEqual(tb[-1 + stack_adj].name, 'run_test')
            else:
                self.assertFalse('source_traceback' in context)

        for flag in (True, False):
            for meth_name, meth, stack_adj in (
                ('call_soon',
                    self.loop.call_soon, 0),
                ('call_later',  # `-1` accounts for lambda
                    lambda *args: self.loop.call_later(0.01, *args), -1)
            ):
                with self.subTest(debug=flag, meth_name=meth_name):
                    run_test(flag, meth, stack_adj)

    def test_now_update(self):
        async def run():
            st = self.loop.time()
            time.sleep(0.05)
            return self.loop.time() - st

        delta = self.loop.run_until_complete(run())
        self.assertTrue(delta > 0.049 and delta < 0.6)

    def test_call_later(self):
        calls = []

        def cb(inc=10, stop=False):
            calls.append(inc)
            self.assertTrue(self.loop.is_running())
            if stop:
                self.loop.call_soon(self.loop.stop)

        self.loop.call_later(0.05, cb)

        # canceled right away
        h = self.loop.call_later(0.05, cb, 100, True)
        self.assertIn('.cb', repr(h))
        h.cancel()
        self.assertIn('cancelled', repr(h))

        self.loop.call_later(0.05, cb, 1, True)
        self.loop.call_later(1000, cb, 1000)  # shouldn't be called

        started = time.monotonic()
        self.loop.run_forever()
        finished = time.monotonic()

        self.assertLess(finished - started, 0.1)
        self.assertGreater(finished - started, 0.04)

        self.assertEqual(calls, [10, 1])

        self.assertFalse(self.loop.is_running())

    def test_call_later_negative(self):
        calls = []

        def cb(arg):
            calls.append(arg)
            self.loop.stop()

        self.loop.call_later(-1, cb, 'a')
        self.loop.run_forever()
        self.assertEqual(calls, ['a'])

    def test_call_at(self):
        i = 0

        def cb(inc):
            nonlocal i
            i += inc
            self.loop.stop()

        at = self.loop.time() + 0.05

        self.loop.call_at(at, cb, 10)
        self.loop.call_at(at, cb, 100).cancel()
        self.loop.call_at(at, cb, 1)

        started = time.monotonic()
        self.loop.run_forever()
        finished = time.monotonic()

        self.assertLess(finished - started, 0.07)
        self.assertGreater(finished - started, 0.045)

        self.assertEqual(i, 11)

    def test_check_thread(self):
        def check_thread(loop, debug):
            def cb():
                pass

            loop.set_debug(debug)
            if debug:
                msg = ("Non-thread-safe operation invoked on an "
                       "event loop other than the current one")
                with self.assertRaisesRegex(RuntimeError, msg):
                    loop.call_soon(cb)
                with self.assertRaisesRegex(RuntimeError, msg):
                    loop.call_later(60, cb)
                with self.assertRaisesRegex(RuntimeError, msg):
                    loop.call_at(loop.time() + 60, cb)
            else:
                loop.call_soon(cb)
                loop.call_later(60, cb)
                loop.call_at(loop.time() + 60, cb)

        def check_in_thread(loop, event, debug, create_loop, fut):
            # wait until the event loop is running
            event.wait()

            try:
                if create_loop:
                    loop2 = self.new_loop()
                    try:
                        asyncio.set_event_loop(loop2)
                        check_thread(loop, debug)
                    finally:
                        asyncio.set_event_loop(None)
                        loop2.close()
                else:
                    check_thread(loop, debug)
            except Exception as exc:
                loop.call_soon_threadsafe(fut.set_exception, exc)
            else:
                loop.call_soon_threadsafe(fut.set_result, None)

        def test_thread(loop, debug, create_loop=False):
            event = threading.Event()
            fut = asyncio.Future(loop=loop)
            loop.call_soon(event.set)
            args = (loop, event, debug, create_loop, fut)
            thread = threading.Thread(target=check_in_thread, args=args)
            thread.start()
            loop.run_until_complete(fut)
            thread.join()

        # raise RuntimeError if the thread has no event loop
        test_thread(self.loop, True)

        # check disabled if debug mode is disabled
        test_thread(self.loop, False)

        # raise RuntimeError if the event loop of the thread is not the called
        # event loop
        test_thread(self.loop, True, create_loop=True)

        # check disabled if debug mode is disabled
        test_thread(self.loop, False, create_loop=True)

    def test_run_once_in_executor_plain(self):
        called = []

        def cb(arg):
            called.append(arg)

        async def runner():
            await self.loop.run_in_executor(None, cb, 'a')

        self.loop.run_until_complete(runner())

        self.assertEqual(called, ['a'])

    def test_set_debug(self):
        self.loop.set_debug(True)
        self.assertTrue(self.loop.get_debug())
        self.loop.set_debug(False)
        self.assertFalse(self.loop.get_debug())

    def test_run_until_complete_type_error(self):
        self.assertRaises(
            TypeError, self.loop.run_until_complete, 'blah')

    def test_run_until_complete_loop(self):
        task = asyncio.Future(loop=self.loop)
        other_loop = self.new_loop()
        self.addCleanup(other_loop.close)
        self.assertRaises(
            ValueError, other_loop.run_until_complete, task)

    def test_run_until_complete_error(self):
        async def foo():
            raise ValueError('aaa')
        with self.assertRaisesRegex(ValueError, 'aaa'):
            self.loop.run_until_complete(foo())

    def test_default_exc_handler_callback(self):
        self.loop._process_events = mock.Mock()

        def zero_error(fut):
            fut.set_result(True)
            1 / 0

        logger = logging.getLogger('asyncio')

        # Test call_soon (events.Handle)
        with mock.patch.object(logger, 'error') as log:
            fut = asyncio.Future(loop=self.loop)
            self.loop.call_soon(zero_error, fut)
            fut.add_done_callback(lambda fut: self.loop.stop())
            self.loop.run_forever()
            log.assert_called_with(
                self.mock_pattern('Exception in callback.*zero'),
                exc_info=mock.ANY)

        # Test call_later (events.TimerHandle)
        with mock.patch.object(logger, 'error') as log:
            fut = asyncio.Future(loop=self.loop)
            self.loop.call_later(0.01, zero_error, fut)
            fut.add_done_callback(lambda fut: self.loop.stop())
            self.loop.run_forever()
            log.assert_called_with(
                self.mock_pattern('Exception in callback.*zero'),
                exc_info=mock.ANY)

    def test_set_exc_handler_custom(self):
        logger = logging.getLogger('asyncio')

        def run_loop():
            def zero_error():
                self.loop.stop()
                1 / 0
            self.loop.call_soon(zero_error)
            self.loop.run_forever()

        errors = []

        def handler(loop, exc):
            errors.append(exc)

        self.loop.set_debug(True)

        if hasattr(self.loop, 'get_exception_handler'):
            # Available since Python 3.5.2
            self.assertIsNone(self.loop.get_exception_handler())
        self.loop.set_exception_handler(handler)
        if hasattr(self.loop, 'get_exception_handler'):
            self.assertIs(self.loop.get_exception_handler(), handler)
        run_loop()
        self.assertEqual(len(errors), 1)
        self.assertRegex(errors[-1]['message'],
                         'Exception in callback.*zero_error')

        self.loop.set_exception_handler(None)
        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                self.mock_pattern('Exception in callback.*zero'),
                exc_info=mock.ANY)

        self.assertEqual(len(errors), 1)

    def test_set_exc_handler_broken(self):
        logger = logging.getLogger('asyncio')

        def run_loop():
            def zero_error():
                self.loop.stop()
                1 / 0
            self.loop.call_soon(zero_error)
            self.loop.run_forever()

        def handler(loop, context):
            raise AttributeError('spam')

        self.loop._process_events = mock.Mock()

        self.loop.set_exception_handler(handler)

        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                self.mock_pattern('Unhandled error in exception handler'),
                exc_info=mock.ANY)

    def test_default_exc_handler_broken(self):
        logger = logging.getLogger('asyncio')
        _context = None

        class Loop(uvloop.Loop):

            _selector = mock.Mock()
            _process_events = mock.Mock()

            def default_exception_handler(self, context):
                nonlocal _context
                _context = context
                # Simulates custom buggy "default_exception_handler"
                raise ValueError('spam')

        loop = Loop()
        self.addCleanup(loop.close)
        asyncio.set_event_loop(loop)

        def run_loop():
            def zero_error():
                loop.stop()
                1 / 0
            loop.call_soon(zero_error)
            loop.run_forever()

        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                'Exception in default exception handler',
                exc_info=True)

        def custom_handler(loop, context):
            raise ValueError('ham')

        _context = None
        loop.set_exception_handler(custom_handler)
        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                self.mock_pattern('Exception in default exception.*'
                                  'while handling.*in custom'),
                exc_info=True)

            # Check that original context was passed to default
            # exception handler.
            self.assertIn('context', _context)
            self.assertIs(type(_context['context']['exception']),
                          ZeroDivisionError)

    def test_set_task_factory_invalid(self):
        with self.assertRaisesRegex(
                TypeError,
                'task factory must be a callable or None'):

            self.loop.set_task_factory(1)

        self.assertIsNone(self.loop.get_task_factory())

    def test_set_task_factory(self):
        self.loop._process_events = mock.Mock()

        class MyTask(asyncio.Task):
            pass

        @asyncio.coroutine
        def coro():
            pass

        factory = lambda loop, coro: MyTask(coro, loop=loop)

        self.assertIsNone(self.loop.get_task_factory())
        self.loop.set_task_factory(factory)
        self.assertIs(self.loop.get_task_factory(), factory)

        task = self.loop.create_task(coro())
        self.assertTrue(isinstance(task, MyTask))
        self.loop.run_until_complete(task)

        self.loop.set_task_factory(None)
        self.assertIsNone(self.loop.get_task_factory())

        task = self.loop.create_task(coro())
        self.assertTrue(isinstance(task, asyncio.Task))
        self.assertFalse(isinstance(task, MyTask))
        self.loop.run_until_complete(task)

class TestBaseUV(_TestBase, UVTestCase):

    def test_loop_create_future(self):
        fut = self.loop.create_future()
        self.assertTrue(isinstance(fut, asyncio.Future))
        self.assertIs(fut._loop, self.loop)
        fut.cancel()

class TestBaseAIO(_TestBase, AIOTestCase):
    pass

class TestPolicy(unittest.TestCase):

    def test_uvloop_policy(self):
        try:
            asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
            loop = asyncio.new_event_loop()
            try:
                self.assertIsInstance(loop, uvloop.Loop)
            finally:
                loop.close()
        finally:
            asyncio.set_event_loop_policy(None)
<commit_msg>tests: Make call_at test more stable<commit_after>import asyncio
import logging
import threading
import time
import uvloop
import unittest
import weakref

from unittest import mock
from uvloop._testbase import UVTestCase, AIOTestCase

class _TestBase:

    def test_close(self):
        self.assertFalse(self.loop._closed)
        self.assertFalse(self.loop.is_closed())
        self.loop.close()
        self.assertTrue(self.loop._closed)
        self.assertTrue(self.loop.is_closed())

        # it should be possible to call close() more than once
        self.loop.close()
        self.loop.close()

        # operation blocked when the loop is closed
        f = asyncio.Future(loop=self.loop)
        self.assertRaises(RuntimeError, self.loop.run_forever)
        self.assertRaises(RuntimeError, self.loop.run_until_complete, f)

    def test_handle_weakref(self):
        wd = weakref.WeakValueDictionary()
        h = self.loop.call_soon(lambda: None)
        wd['h'] = h  # Would fail without __weakref__ slot.

    def test_call_soon(self):
        calls = []

        def cb(inc):
            calls.append(inc)
            self.loop.stop()

        self.loop.call_soon(cb, 10)

        h = self.loop.call_soon(cb, 100)
        self.assertIn('.cb', repr(h))
        h.cancel()
        self.assertIn('cancelled', repr(h))

        self.loop.call_soon(cb, 1)

        self.loop.run_forever()

        self.assertEqual(calls, [10, 1])

    def test_call_soon_base_exc(self):
        def cb():
            raise KeyboardInterrupt()

        self.loop.call_soon(cb)

        with self.assertRaises(KeyboardInterrupt):
            self.loop.run_forever()

        self.assertFalse(self.loop.is_closed())

    def test_calls_debug_reporting(self):
        def run_test(debug, meth, stack_adj):
            context = None

            def handler(loop, ctx):
                nonlocal context
                context = ctx

            self.loop.set_debug(debug)
            self.loop.set_exception_handler(handler)

            def cb():
                1 / 0

            meth(cb)
            self.assertIsNone(context)
            self.loop.run_until_complete(asyncio.sleep(0.05, loop=self.loop))

            self.assertIs(type(context['exception']), ZeroDivisionError)
            self.assertTrue(context['message'].startswith(
                'Exception in callback'))

            if debug:
                tb = context['source_traceback']
                self.assertEqual(tb[-1 + stack_adj].name, 'run_test')
            else:
                self.assertFalse('source_traceback' in context)

        for flag in (True, False):
            for meth_name, meth, stack_adj in (
                ('call_soon',
                    self.loop.call_soon, 0),
                ('call_later',  # `-1` accounts for lambda
                    lambda *args: self.loop.call_later(0.01, *args), -1)
            ):
                with self.subTest(debug=flag, meth_name=meth_name):
                    run_test(flag, meth, stack_adj)

    def test_now_update(self):
        async def run():
            st = self.loop.time()
            time.sleep(0.05)
            return self.loop.time() - st

        delta = self.loop.run_until_complete(run())
        self.assertTrue(delta > 0.049 and delta < 0.6)

    def test_call_later(self):
        calls = []

        def cb(inc=10, stop=False):
            calls.append(inc)
            self.assertTrue(self.loop.is_running())
            if stop:
                self.loop.call_soon(self.loop.stop)

        self.loop.call_later(0.05, cb)

        # canceled right away
        h = self.loop.call_later(0.05, cb, 100, True)
        self.assertIn('.cb', repr(h))
        h.cancel()
        self.assertIn('cancelled', repr(h))

        self.loop.call_later(0.05, cb, 1, True)
        self.loop.call_later(1000, cb, 1000)  # shouldn't be called

        started = time.monotonic()
        self.loop.run_forever()
        finished = time.monotonic()

        self.assertLess(finished - started, 0.1)
        self.assertGreater(finished - started, 0.04)

        self.assertEqual(calls, [10, 1])

        self.assertFalse(self.loop.is_running())

    def test_call_later_negative(self):
        calls = []

        def cb(arg):
            calls.append(arg)
            self.loop.stop()

        self.loop.call_later(-1, cb, 'a')
        self.loop.run_forever()
        self.assertEqual(calls, ['a'])

    def test_call_at(self):
        i = 0

        def cb(inc):
            nonlocal i
            i += inc
            self.loop.stop()

        at = self.loop.time() + 0.05

        self.loop.call_at(at, cb, 100).cancel()
        self.loop.call_at(at, cb, 10)

        started = time.monotonic()
        self.loop.run_forever()
        finished = time.monotonic()

        self.assertLess(finished - started, 0.07)
        self.assertGreater(finished - started, 0.045)

        self.assertEqual(i, 10)

    def test_check_thread(self):
        def check_thread(loop, debug):
            def cb():
                pass

            loop.set_debug(debug)
            if debug:
                msg = ("Non-thread-safe operation invoked on an "
                       "event loop other than the current one")
                with self.assertRaisesRegex(RuntimeError, msg):
                    loop.call_soon(cb)
                with self.assertRaisesRegex(RuntimeError, msg):
                    loop.call_later(60, cb)
                with self.assertRaisesRegex(RuntimeError, msg):
                    loop.call_at(loop.time() + 60, cb)
            else:
                loop.call_soon(cb)
                loop.call_later(60, cb)
                loop.call_at(loop.time() + 60, cb)

        def check_in_thread(loop, event, debug, create_loop, fut):
            # wait until the event loop is running
            event.wait()

            try:
                if create_loop:
                    loop2 = self.new_loop()
                    try:
                        asyncio.set_event_loop(loop2)
                        check_thread(loop, debug)
                    finally:
                        asyncio.set_event_loop(None)
                        loop2.close()
                else:
                    check_thread(loop, debug)
            except Exception as exc:
                loop.call_soon_threadsafe(fut.set_exception, exc)
            else:
                loop.call_soon_threadsafe(fut.set_result, None)

        def test_thread(loop, debug, create_loop=False):
            event = threading.Event()
            fut = asyncio.Future(loop=loop)
            loop.call_soon(event.set)
            args = (loop, event, debug, create_loop, fut)
            thread = threading.Thread(target=check_in_thread, args=args)
            thread.start()
            loop.run_until_complete(fut)
            thread.join()

        # raise RuntimeError if the thread has no event loop
        test_thread(self.loop, True)

        # check disabled if debug mode is disabled
        test_thread(self.loop, False)

        # raise RuntimeError if the event loop of the thread is not the called
        # event loop
        test_thread(self.loop, True, create_loop=True)

        # check disabled if debug mode is disabled
        test_thread(self.loop, False, create_loop=True)

    def test_run_once_in_executor_plain(self):
        called = []

        def cb(arg):
            called.append(arg)

        async def runner():
            await self.loop.run_in_executor(None, cb, 'a')

        self.loop.run_until_complete(runner())

        self.assertEqual(called, ['a'])

    def test_set_debug(self):
        self.loop.set_debug(True)
        self.assertTrue(self.loop.get_debug())
        self.loop.set_debug(False)
        self.assertFalse(self.loop.get_debug())

    def test_run_until_complete_type_error(self):
        self.assertRaises(
            TypeError, self.loop.run_until_complete, 'blah')

    def test_run_until_complete_loop(self):
        task = asyncio.Future(loop=self.loop)
        other_loop = self.new_loop()
        self.addCleanup(other_loop.close)
        self.assertRaises(
            ValueError, other_loop.run_until_complete, task)

    def test_run_until_complete_error(self):
        async def foo():
            raise ValueError('aaa')
        with self.assertRaisesRegex(ValueError, 'aaa'):
            self.loop.run_until_complete(foo())

    def test_default_exc_handler_callback(self):
        self.loop._process_events = mock.Mock()

        def zero_error(fut):
            fut.set_result(True)
            1 / 0

        logger = logging.getLogger('asyncio')

        # Test call_soon (events.Handle)
        with mock.patch.object(logger, 'error') as log:
            fut = asyncio.Future(loop=self.loop)
            self.loop.call_soon(zero_error, fut)
            fut.add_done_callback(lambda fut: self.loop.stop())
            self.loop.run_forever()
            log.assert_called_with(
                self.mock_pattern('Exception in callback.*zero'),
                exc_info=mock.ANY)

        # Test call_later (events.TimerHandle)
        with mock.patch.object(logger, 'error') as log:
            fut = asyncio.Future(loop=self.loop)
            self.loop.call_later(0.01, zero_error, fut)
            fut.add_done_callback(lambda fut: self.loop.stop())
            self.loop.run_forever()
            log.assert_called_with(
                self.mock_pattern('Exception in callback.*zero'),
                exc_info=mock.ANY)

    def test_set_exc_handler_custom(self):
        logger = logging.getLogger('asyncio')

        def run_loop():
            def zero_error():
                self.loop.stop()
                1 / 0
            self.loop.call_soon(zero_error)
            self.loop.run_forever()

        errors = []

        def handler(loop, exc):
            errors.append(exc)

        self.loop.set_debug(True)

        if hasattr(self.loop, 'get_exception_handler'):
            # Available since Python 3.5.2
            self.assertIsNone(self.loop.get_exception_handler())
        self.loop.set_exception_handler(handler)
        if hasattr(self.loop, 'get_exception_handler'):
            self.assertIs(self.loop.get_exception_handler(), handler)
        run_loop()
        self.assertEqual(len(errors), 1)
        self.assertRegex(errors[-1]['message'],
                         'Exception in callback.*zero_error')

        self.loop.set_exception_handler(None)
        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                self.mock_pattern('Exception in callback.*zero'),
                exc_info=mock.ANY)

        self.assertEqual(len(errors), 1)

    def test_set_exc_handler_broken(self):
        logger = logging.getLogger('asyncio')

        def run_loop():
            def zero_error():
                self.loop.stop()
                1 / 0
            self.loop.call_soon(zero_error)
            self.loop.run_forever()

        def handler(loop, context):
            raise AttributeError('spam')

        self.loop._process_events = mock.Mock()

        self.loop.set_exception_handler(handler)

        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                self.mock_pattern('Unhandled error in exception handler'),
                exc_info=mock.ANY)

    def test_default_exc_handler_broken(self):
        logger = logging.getLogger('asyncio')
        _context = None

        class Loop(uvloop.Loop):

            _selector = mock.Mock()
            _process_events = mock.Mock()

            def default_exception_handler(self, context):
                nonlocal _context
                _context = context
                # Simulates custom buggy "default_exception_handler"
                raise ValueError('spam')

        loop = Loop()
        self.addCleanup(loop.close)
        asyncio.set_event_loop(loop)

        def run_loop():
            def zero_error():
                loop.stop()
                1 / 0
            loop.call_soon(zero_error)
            loop.run_forever()

        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                'Exception in default exception handler',
                exc_info=True)

        def custom_handler(loop, context):
            raise ValueError('ham')

        _context = None
        loop.set_exception_handler(custom_handler)
        with mock.patch.object(logger, 'error') as log:
            run_loop()
            log.assert_called_with(
                self.mock_pattern('Exception in default exception.*'
                                  'while handling.*in custom'),
                exc_info=True)

            # Check that original context was passed to default
            # exception handler.
            self.assertIn('context', _context)
            self.assertIs(type(_context['context']['exception']),
                          ZeroDivisionError)

    def test_set_task_factory_invalid(self):
        with self.assertRaisesRegex(
                TypeError,
                'task factory must be a callable or None'):

            self.loop.set_task_factory(1)

        self.assertIsNone(self.loop.get_task_factory())

    def test_set_task_factory(self):
        self.loop._process_events = mock.Mock()

        class MyTask(asyncio.Task):
            pass

        @asyncio.coroutine
        def coro():
            pass

        factory = lambda loop, coro: MyTask(coro, loop=loop)

        self.assertIsNone(self.loop.get_task_factory())
        self.loop.set_task_factory(factory)
        self.assertIs(self.loop.get_task_factory(), factory)

        task = self.loop.create_task(coro())
        self.assertTrue(isinstance(task, MyTask))
        self.loop.run_until_complete(task)

        self.loop.set_task_factory(None)
        self.assertIsNone(self.loop.get_task_factory())

        task = self.loop.create_task(coro())
        self.assertTrue(isinstance(task, asyncio.Task))
        self.assertFalse(isinstance(task, MyTask))
        self.loop.run_until_complete(task)

class TestBaseUV(_TestBase, UVTestCase):

    def test_loop_create_future(self):
        fut = self.loop.create_future()
        self.assertTrue(isinstance(fut, asyncio.Future))
        self.assertIs(fut._loop, self.loop)
        fut.cancel()

class TestBaseAIO(_TestBase, AIOTestCase):
    pass

class TestPolicy(unittest.TestCase):

    def test_uvloop_policy(self):
        try:
            asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
            loop = asyncio.new_event_loop()
            try:
                self.assertIsInstance(loop, uvloop.Loop)
            finally:
                loop.close()
        finally:
            asyncio.set_event_loop_policy(None)
<|endoftext|>


<commit_before>import datajoint as dj
import timeit
import numpy as np
import uuid
from faker import Faker
from . import schema
from decimal import Decimal
from datetime import datetime
from datajoint.blob import pack, unpack
from numpy.testing import assert_array_equal
from nose.tools import (
    assert_equal,
    assert_true,
    assert_false,
    assert_list_equal,
    assert_set_equal,
    assert_tuple_equal,
    assert_dict_equal,
)

def test_pack():

    for x in (
        32,
        -3.7e-2,
        np.float64(3e31),
        -np.inf,
        np.int8(-3),
        np.uint8(-1),
        np.int16(-33),
        np.uint16(-33),
        np.int32(-3),
        np.uint32(-1),
        np.int64(373),
        np.uint64(-3),
    ):
        assert_equal(x, unpack(pack(x)), "Scalars don't match!")

    x = np.nan
    assert_true(np.isnan(unpack(pack(x))), "nan scalar did not match!")

    x = np.random.randn(8, 10)
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = np.random.randn(10)
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = 7j
    assert_equal(x, unpack(pack(x)), "Complex scalar does not match")

    x = np.float32(np.random.randn(3, 4, 5))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = np.int16(np.random.randn(1, 2, 3))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = None
    assert_true(unpack(pack(x)) is None, "None did not match")

    x = -255
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, int) and not isinstance(y, np.ndarray),
        "Scalar int did not match",
    )

    x = -25523987234234287910987234987098245697129798713407812347
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, int) and not isinstance(y, np.ndarray),
        "Unbounded int did not match",
    )

    x = 7.0
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, float) and not isinstance(y, np.ndarray),
        "Scalar float did not match",
    )

    x = 7j
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, complex) and not isinstance(y, np.ndarray),
        "Complex scalar did not match",
    )

    x = True
    assert_true(unpack(pack(x)) is True, "Scalar bool did not match")

    x = [None]
    assert_list_equal(x, unpack(pack(x)))

    x = {
        "name": "Anonymous",
        "age": 15,
        99: datetime.now(),
        "range": [110, 190],
        (11, 12): None,
    }
    y = unpack(pack(x))
    assert_dict_equal(x, y, "Dict do not match!")
    assert_false(
        isinstance(["range"][0], np.ndarray), "Scalar int was coerced into arrray."
    )

    x = uuid.uuid4()
    assert_equal(x, unpack(pack(x)), "UUID did not match")

    x = Decimal("-112122121.000003000")
    assert_equal(x, unpack(pack(x)), "Decimal did not pack/unpack correctly")

    x = [1, datetime.now(), {1: "one", "two": 2}, (1, 2)]
    assert_list_equal(x, unpack(pack(x)), "List did not pack/unpack correctly")

    x = (1, datetime.now(), {1: "one", "two": 2}, (uuid.uuid4(), 2))
    assert_tuple_equal(x, unpack(pack(x)), "Tuple did not pack/unpack correctly")

    x = (
        1,
        {datetime.now().date(): "today", "now": datetime.now().date()},
        {"yes!": [1, 2, np.array((3, 4))]},
    )
    y = unpack(pack(x))
    assert_dict_equal(x[1], y[1])
    assert_array_equal(x[2]["yes!"][2], y[2]["yes!"][2])

    x = {"elephant"}
    assert_set_equal(x, unpack(pack(x)), "Set did not pack/unpack correctly")

    x = tuple(range(10))
    assert_tuple_equal(
        x, unpack(pack(range(10))), "Iterator did not pack/unpack correctly"
    )

    x = Decimal("1.24")
    assert_true(x == unpack(pack(x)), "Decimal object did not pack/unpack correctly")

    x = datetime.now()
    assert_true(x == unpack(pack(x)), "Datetime object did not pack/unpack correctly")

    x = np.bool_(True)
    assert_true(x == unpack(pack(x)), "Numpy bool object did not pack/unpack correctly")

    x = "test"
    assert_true(x == unpack(pack(x)), "String object did not pack/unpack correctly")

    x = np.array(["yes"])
    assert_true(
        x == unpack(pack(x)), "Numpy string array object did not pack/unpack correctly"
    )

    x = np.datetime64("1998").astype("datetime64[us]")
    assert_true(x == unpack(pack(x)))

def test_recarrays():
    x = np.array([(1.0, 2), (3.0, 4)], dtype=[("x", float), ("y", int)])
    assert_array_equal(x, unpack(pack(x)))

    x = x.view(np.recarray)
    assert_array_equal(x, unpack(pack(x)))

    x = np.array([(3, 4)], dtype=[("tmp0", float), ("tmp1", "O")]).view(np.recarray)
    assert_array_equal(x, unpack(pack(x)))

def test_object_arrays():
    x = np.array(((1, 2, 3), True))
    assert_array_equal(x, unpack(pack(x)), "Object array did not serialize correctly")

def test_complex():
    z = np.random.randn(8, 10) + 1j * np.random.randn(8, 10)
    assert_array_equal(z, unpack(pack(z)), "Arrays do not match!")

    z = np.random.randn(10) + 1j * np.random.randn(10)
    assert_array_equal(z, unpack(pack(z)), "Arrays do not match!")

    x = np.float32(np.random.randn(3, 4, 5)) + 1j * np.float32(np.random.randn(3, 4, 5))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = np.int16(np.random.randn(1, 2, 3)) + 1j * np.int16(np.random.randn(1, 2, 3))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

def test_insert_longblob():
    insert_dj_blob = {"id": 1, "data": [1, 2, 3]}
    schema.Longblob.insert1(insert_dj_blob)
    assert (schema.Longblob & "id=1").fetch1() == insert_dj_blob
    (schema.Longblob & "id=1").delete()

    query_mym_blob = {"id": 1, "data": np.array([1, 2, 3])}
    schema.Longblob.insert1(query_mym_blob)
    assert (schema.Longblob & "id=1").fetch1()["data"].all() == query_mym_blob[
        "data"
    ].all()
    (schema.Longblob & "id=1").delete()

    query_32_blob = (
        "INSERT INTO djtest_test1.longblob (id, data) VALUES (1, "
        "X'6D596D00530200000001000000010000000400000068697473007369646573007461736B73007374"
        "616765004D000000410200000001000000070000000600000000000000000000000000F8FF00000000"
        "0000F03F000000000000F03F0000000000000000000000000000F03F00000000000000000000000000"
        "00F8FF230000004102000000010000000700000004000000000000006C006C006C006C00720072006C"
        "0023000000410200000001000000070000000400000000000000640064006400640064006400640025"
        "00000041020000000100000008000000040000000000000053007400610067006500200031003000')"
    )
    dj.conn().query(query_32_blob).fetchall()
    dj.blob.use_32bit_dims = True
    assert (schema.Longblob & "id=1").fetch1() == {
        "id": 1,
        "data": np.rec.array(
            [
                [
                    (
                        np.array([[np.nan, 1.0, 1.0, 0.0, 1.0, 0.0, np.nan]]),
                        np.array(["llllrrl"], dtype="<U7"),
                        np.array(["ddddddd"], dtype="<U7"),
                        np.array(["Stage 10"], dtype="<U8"),
                    )
                ]
            ],
            dtype=[("hits", "O"), ("sides", "O"), ("tasks", "O"), ("stage", "O")],
        ),
    }
    (schema.Longblob & "id=1").delete()
    dj.blob.use_32bit_dims = False

def test_datetime_serialization_speed():
    # If this fails that means for some reason deserializing/serializing
    # np arrays of np.datetime64 types is now slower than regular arrays of datetime64

    optimized_exe_time = timeit.timeit(
        setup="myarr=pack(np.array([np.datetime64('2022-10-13 03:03:13') for _ in range(0, 10000)]))",
        stmt="unpack(myarr)",
        number=10,
        globals=globals(),
    )
    baseline_exe_time = timeit.timeit(
        setup="myarr2=pack(np.array([datetime(2022,10,13,3,3,13) for _ in range (0, 10000)]))",
        stmt="unpack(myarr2)",
        number=10,
        globals=globals(),
    )

    assert optimized_exe_time * 1000 < baseline_exe_time
<commit_msg>small changes.<commit_after>import datajoint as dj
import timeit
import numpy as np
import uuid
from faker import Faker
from . import schema
from decimal import Decimal
from datetime import datetime
from datajoint.blob import pack, unpack
from numpy.testing import assert_array_equal
from nose.tools import (
    assert_equal,
    assert_true,
    assert_false,
    assert_list_equal,
    assert_set_equal,
    assert_tuple_equal,
    assert_dict_equal,
)

def test_pack():

    for x in (
        32,
        -3.7e-2,
        np.float64(3e31),
        -np.inf,
        np.int8(-3),
        np.uint8(-1),
        np.int16(-33),
        np.uint16(-33),
        np.int32(-3),
        np.uint32(-1),
        np.int64(373),
        np.uint64(-3),
    ):
        assert_equal(x, unpack(pack(x)), "Scalars don't match!")

    x = np.nan
    assert_true(np.isnan(unpack(pack(x))), "nan scalar did not match!")

    x = np.random.randn(8, 10)
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = np.random.randn(10)
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = 7j
    assert_equal(x, unpack(pack(x)), "Complex scalar does not match")

    x = np.float32(np.random.randn(3, 4, 5))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = np.int16(np.random.randn(1, 2, 3))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = None
    assert_true(unpack(pack(x)) is None, "None did not match")

    x = -255
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, int) and not isinstance(y, np.ndarray),
        "Scalar int did not match",
    )

    x = -25523987234234287910987234987098245697129798713407812347
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, int) and not isinstance(y, np.ndarray),
        "Unbounded int did not match",
    )

    x = 7.0
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, float) and not isinstance(y, np.ndarray),
        "Scalar float did not match",
    )

    x = 7j
    y = unpack(pack(x))
    assert_true(
        x == y and isinstance(y, complex) and not isinstance(y, np.ndarray),
        "Complex scalar did not match",
    )

    x = True
    assert_true(unpack(pack(x)) is True, "Scalar bool did not match")

    x = [None]
    assert_list_equal(x, unpack(pack(x)))

    x = {
        "name": "Anonymous",
        "age": 15,
        99: datetime.now(),
        "range": [110, 190],
        (11, 12): None,
    }
    y = unpack(pack(x))
    assert_dict_equal(x, y, "Dict do not match!")
    assert_false(
        isinstance(["range"][0], np.ndarray), "Scalar int was coerced into arrray."
    )

    x = uuid.uuid4()
    assert_equal(x, unpack(pack(x)), "UUID did not match")

    x = Decimal("-112122121.000003000")
    assert_equal(x, unpack(pack(x)), "Decimal did not pack/unpack correctly")

    x = [1, datetime.now(), {1: "one", "two": 2}, (1, 2)]
    assert_list_equal(x, unpack(pack(x)), "List did not pack/unpack correctly")

    x = (1, datetime.now(), {1: "one", "two": 2}, (uuid.uuid4(), 2))
    assert_tuple_equal(x, unpack(pack(x)), "Tuple did not pack/unpack correctly")

    x = (
        1,
        {datetime.now().date(): "today", "now": datetime.now().date()},
        {"yes!": [1, 2, np.array((3, 4))]},
    )
    y = unpack(pack(x))
    assert_dict_equal(x[1], y[1])
    assert_array_equal(x[2]["yes!"][2], y[2]["yes!"][2])

    x = {"elephant"}
    assert_set_equal(x, unpack(pack(x)), "Set did not pack/unpack correctly")

    x = tuple(range(10))
    assert_tuple_equal(
        x, unpack(pack(range(10))), "Iterator did not pack/unpack correctly"
    )

    x = Decimal("1.24")
    assert_true(x == unpack(pack(x)), "Decimal object did not pack/unpack correctly")

    x = datetime.now()
    assert_true(x == unpack(pack(x)), "Datetime object did not pack/unpack correctly")

    x = np.bool_(True)
    assert_true(x == unpack(pack(x)), "Numpy bool object did not pack/unpack correctly")

    x = "test"
    assert_true(x == unpack(pack(x)), "String object did not pack/unpack correctly")

    x = np.array(["yes"])
    assert_true(
        x == unpack(pack(x)), "Numpy string array object did not pack/unpack correctly"
    )

    x = np.datetime64("1998").astype("datetime64[us]")
    assert_true(x == unpack(pack(x)))

def test_recarrays():
    x = np.array([(1.0, 2), (3.0, 4)], dtype=[("x", float), ("y", int)])
    assert_array_equal(x, unpack(pack(x)))

    x = x.view(np.recarray)
    assert_array_equal(x, unpack(pack(x)))

    x = np.array([(3, 4)], dtype=[("tmp0", float), ("tmp1", "O")]).view(np.recarray)
    assert_array_equal(x, unpack(pack(x)))

def test_object_arrays():
    x = np.array(((1, 2, 3), True))
    assert_array_equal(x, unpack(pack(x)), "Object array did not serialize correctly")

def test_complex():
    z = np.random.randn(8, 10) + 1j * np.random.randn(8, 10)
    assert_array_equal(z, unpack(pack(z)), "Arrays do not match!")

    z = np.random.randn(10) + 1j * np.random.randn(10)
    assert_array_equal(z, unpack(pack(z)), "Arrays do not match!")

    x = np.float32(np.random.randn(3, 4, 5)) + 1j * np.float32(np.random.randn(3, 4, 5))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

    x = np.int16(np.random.randn(1, 2, 3)) + 1j * np.int16(np.random.randn(1, 2, 3))
    assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")

def test_insert_longblob():
    insert_dj_blob = {"id": 1, "data": [1, 2, 3]}
    schema.Longblob.insert1(insert_dj_blob)
    assert (schema.Longblob & "id=1").fetch1() == insert_dj_blob
    (schema.Longblob & "id=1").delete()

    query_mym_blob = {"id": 1, "data": np.array([1, 2, 3])}
    schema.Longblob.insert1(query_mym_blob)
    assert (schema.Longblob & "id=1").fetch1()["data"].all() == query_mym_blob[
        "data"
    ].all()
    (schema.Longblob & "id=1").delete()

    query_32_blob = (
        "INSERT INTO djtest_test1.longblob (id, data) VALUES (1, "
        "X'6D596D00530200000001000000010000000400000068697473007369646573007461736B73007374"
        "616765004D000000410200000001000000070000000600000000000000000000000000F8FF00000000"
        "0000F03F000000000000F03F0000000000000000000000000000F03F00000000000000000000000000"
        "00F8FF230000004102000000010000000700000004000000000000006C006C006C006C00720072006C"
        "0023000000410200000001000000070000000400000000000000640064006400640064006400640025"
        "00000041020000000100000008000000040000000000000053007400610067006500200031003000')"
    )
    dj.conn().query(query_32_blob).fetchall()
    dj.blob.use_32bit_dims = True
    assert (schema.Longblob & "id=1").fetch1() == {
        "id": 1,
        "data": np.rec.array(
            [
                [
                    (
                        np.array([[np.nan, 1.0, 1.0, 0.0, 1.0, 0.0, np.nan]]),
                        np.array(["llllrrl"], dtype="<U7"),
                        np.array(["ddddddd"], dtype="<U7"),
                        np.array(["Stage 10"], dtype="<U8"),
                    )
                ]
            ],
            dtype=[("hits", "O"), ("sides", "O"), ("tasks", "O"), ("stage", "O")],
        ),
    }
    (schema.Longblob & "id=1").delete()
    dj.blob.use_32bit_dims = False

def test_datetime_serialization_speed():
    # If this fails that means for some reason deserializing/serializing
    # np arrays of np.datetime64 types is now slower than regular arrays of datetime64

    optimized_exe_time = timeit.timeit(
        setup="myarr=pack(np.array([np.datetime64('2022-10-13 03:03:13') for _ in range(0, 10000)]))",
        stmt="unpack(myarr)",
        number=10,
        globals=globals(),
    )
    print(f"np time {optimized_exe_time}")
    baseline_exe_time = timeit.timeit(
        setup="myarr2=pack(np.array([datetime(2022,10,13,3,3,13) for _ in range (0, 10000)]))",
        stmt="unpack(myarr2)",
        number=10,
        globals=globals(),
    )
    print(f"python time {baseline_exe_time}")

    assert optimized_exe_time * 1000 < baseline_exe_time
<|endoftext|>


<commit_before>import os.path as osp
import unittest

from celery import Celery
from celery.result import AsyncResult

import yaml

import backache

celery = Celery()

class TestBulkOperation(unittest.TestCase):
    @classmethod
    def _backache(cls, with_celery=False):
        path = osp.splitext(__file__)[0] + '.yml'
        with open(path) as istr:
            config = yaml.load(istr)
        config['operations'] = {
            'op1': cls._operation1,
            'op2': cls._operation2,
        }
        config['celery_app'] = celery
        celery.config_from_object(config['celery'])
        if with_celery:
            return backache.celerize(**config)
        else:
            return backache.Backache(**config)

    def setUp(self):
        self._cache_hits_received = None
        b = self._backache()
        b._config.cache.clear()
        for i in range(2):
            for j in range(2):
                b._config.resource.delete('op-%s' % i, 'key-%s' % i)

    def test_without_celery(self):
        commands = {
            ('op1', 'key1'): {'cb_args': ('arg1', 'arg2')},
            ('op2', 'key2'): {'cb_args': ('arg3', 'arg4')},
        }
        b = self._backache()
        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(self._cache_hits_received, {})
        self.assertItemsEqual(
            cache_misses,
            [('op1', 'key1'), ('op2', 'key2')]
        )
        result, cb_args = b.consume('op1', 'key1')
        self.assertEquals(result, 'op1-key1')
        self.assertItemsEqual(cb_args, ['arg1', 'arg2'])

        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(
            cache_misses,
            [('op2', 'key2')]
        )
        self.assertEqual(
            self._cache_hits_received,
            {
                ('op1', 'key1'): {
                    'cb_args': ('arg1', 'arg2'),
                    'result': 'op1-key1',
                },
            }
        )

    def test_with_celery(self):
        commands = {
            ('op1', 'key1'): {'cb_args': ('arg1', 'arg2')},
            ('op2', 'key2'): {'cb_args': ('arg3', 'arg4')},
        }
        b = self._backache(with_celery=True)
        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(self._cache_hits_received, {})
        results = []
        for cache_miss in cache_misses:
            self.assertIsInstance(cache_miss, AsyncResult)
            results.append(cache_miss.get())
        self.assertEqual(len(results), 2)
        seen_1st_result = seen_2nd_result = 0
        for result, cb_args in results:
            if result == 'op1-key1':
                self.assertItemsEqual(cb_args, ['arg1', 'arg2'])
                seen_1st_result += 1
            elif result == 'op2-key2':
                self.assertItemsEqual(cb_args, ['arg3', 'arg4'])
                seen_2nd_result += 1
        self.assertEqual(seen_1st_result, 1)
        self.assertEqual(seen_2nd_result, 1)

        commands[('op1', 'key2')] = {'cb_args': ('arg5',)}
        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(
            self._cache_hits_received,
            {
                ('op1', 'key1'): {
                    'cb_args': ('arg1', 'arg2'),
                    'result': 'op1-key1',
                },
                ('op2', 'key2'): {
                    'cb_args': ('arg3', 'arg4'),
                    'result': 'op2-key2',
                },
            }
        )
        self.assertEqual(len(cache_misses), 1)
        self.assertIsInstance(cache_misses[0], AsyncResult)
        self.assertEqual(cache_misses[0].get(), ('op1-key2', ['arg5']))

    def _cache_hits_cb(self, cache_hits):
        self._cache_hits_received = cache_hits

    @classmethod
    def _operation1(cls, uri, context):
        return 'op1-%s' % uri

    @classmethod
    def _operation2(cls, uri, context):
        return 'op2-%s' % uri

if __name__ == '__main__':
    unittest.main()
<commit_msg>Fix UT cleanup<commit_after>import os.path as osp
import unittest

from celery import Celery
from celery.result import AsyncResult

import yaml

import backache

celery = Celery()

class TestBulkOperation(unittest.TestCase):
    @classmethod
    def _backache(cls, with_celery=False):
        path = osp.splitext(__file__)[0] + '.yml'
        with open(path) as istr:
            config = yaml.load(istr)
        config['operations'] = {
            'op1': cls._operation1,
            'op2': cls._operation2,
        }
        config['celery_app'] = celery
        celery.config_from_object(config['celery'])
        if with_celery:
            return backache.celerize(**config)
        else:
            return backache.Backache(**config)

    def setUp(self):
        self._cache_hits_received = None
        b = self._backache()
        b._config.cache.clear()
        for i in range(1, 3):
            for j in range(1, 3):
                b._config.resource.delete('op%s' % i, 'key%s' % j)

    def test_without_celery(self):
        commands = {
            ('op1', 'key1'): {'cb_args': ('arg1', 'arg2')},
            ('op2', 'key2'): {'cb_args': ('arg3', 'arg4')},
        }
        b = self._backache()
        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(self._cache_hits_received, {})
        self.assertItemsEqual(
            cache_misses,
            [('op1', 'key1'), ('op2', 'key2')]
        )
        result, cb_args = b.consume('op1', 'key1')
        self.assertEquals(result, 'op1-key1')
        self.assertItemsEqual(cb_args, ['arg1', 'arg2'])

        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(
            cache_misses,
            [('op2', 'key2')]
        )
        self.assertEqual(
            self._cache_hits_received,
            {
                ('op1', 'key1'): {
                    'cb_args': ('arg1', 'arg2'),
                    'result': 'op1-key1',
                },
            }
        )

    def test_with_celery(self):
        commands = {
            ('op1', 'key1'): {'cb_args': ('arg1', 'arg2')},
            ('op2', 'key2'): {'cb_args': ('arg3', 'arg4')},
        }
        b = self._backache(with_celery=True)
        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(self._cache_hits_received, {})
        results = []
        for cache_miss in cache_misses:
            self.assertIsInstance(cache_miss, AsyncResult)
            results.append(cache_miss.get())
        self.assertEqual(len(results), 2)
        seen_1st_result = seen_2nd_result = 0
        for result, cb_args in results:
            if result == 'op1-key1':
                self.assertItemsEqual(cb_args, ['arg1', 'arg2'])
                seen_1st_result += 1
            elif result == 'op2-key2':
                self.assertItemsEqual(cb_args, ['arg3', 'arg4'])
                seen_2nd_result += 1
        self.assertEqual(seen_1st_result, 1)
        self.assertEqual(seen_2nd_result, 1)

        commands[('op1', 'key2')] = {'cb_args': ('arg5',)}
        cache_misses = b.bulk_get_or_delegate(commands, self._cache_hits_cb)
        self.assertEqual(
            self._cache_hits_received,
            {
                ('op1', 'key1'): {
                    'cb_args': ('arg1', 'arg2'),
                    'result': 'op1-key1',
                },
                ('op2', 'key2'): {
                    'cb_args': ('arg3', 'arg4'),
                    'result': 'op2-key2',
                },
            }
        )
        self.assertEqual(len(cache_misses), 1)
        self.assertIsInstance(cache_misses[0], AsyncResult)
        self.assertEqual(cache_misses[0].get(), ('op1-key2', ['arg5']))

    def _cache_hits_cb(self, cache_hits):
        self._cache_hits_received = cache_hits

    @classmethod
    def _operation1(cls, uri, context):
        return 'op1-%s' % uri

    @classmethod
    def _operation2(cls, uri, context):
        return 'op2-%s' % uri

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


<commit_before>from math import sin, cos
import pytest

from funcy.calc import *

def test_memoize():
    @memoize
    def inc(x):
        calls.append(x)
        return x + 1

    calls = []
    assert inc(0) == 1
    assert inc(1) == 2
    assert inc(0) == 1
    assert calls == [0, 1]

    # using kwargs
    assert inc(x=0) == 1
    assert inc(x=1) == 2
    assert inc(x=0) == 1
    assert calls == [0, 1, 0, 1]

def test_memoize_args_kwargs():
    @memoize
    def mul(x, by=1):
        calls.append((x, by))
        return x * by

    calls = []
    assert mul(0) == 0
    assert mul(1) == 1
    assert mul(0) == 0
    assert calls == [(0, 1), (1, 1)]

    # more with kwargs
    assert mul(0, 1) == 0
    assert mul(1, 1) == 1
    assert mul(0, 1) == 0
    assert calls == [(0, 1), (1, 1), (0, 1), (1, 1)]

def test_memoize_memory():
    @memoize
    def inc(x):
        calls.append(x)
        return x + 1

    calls = []
    inc(0)
    inc.memory.clear()
    inc(0)
    assert calls == [0, 0]

def test_memoize_key_func():
    @memoize(key_func=len)
    def inc(s):
        calls.append(s)
        return s * 2

    calls = []
    assert inc('a') == 'aa'
    assert inc('b') == 'aa'
    inc('ab')
    assert calls == ['a', 'ab']

def test_make_lookuper():
    @make_lookuper
    def letter_index():
        return ((c, i) for i, c in enumerate('abcdefghij'))

    assert letter_index('c') == 2
    with pytest.raises(LookupError): letter_index('_')

def test_make_lookuper_nested():
    tables_built = [0]

    @make_lookuper
    def function_table(f):
        tables_built[0] += 1
        return ((x, f(x)) for x in range(10))

    assert function_table(sin)(5) == sin(5)
    assert function_table(cos)(3) == cos(3)
    assert function_table(sin)(3) == sin(3)
    assert tables_built[0] == 2

    with pytest.raises(LookupError): function_table(cos)(-1)

def test_silent_lookuper():
    @silent_lookuper
    def letter_index():
        return ((c, i) for i, c in enumerate('abcdefghij'))

    assert letter_index('c') == 2
    assert letter_index('_') is None

def test_silnent_lookuper_nested():
    @silent_lookuper
    def function_table(f):
        return ((x, f(x)) for x in range(10))

    assert function_table(sin)(5) == sin(5)
    assert function_table(cos)(-1) is None

def test_cache():
    calls = []

    @cache(timeout=60)
    def inc(x):
        calls.append(x)
        return x + 1

    assert inc(0) == 1
    assert inc(1) == 2
    assert inc(0) == 1
    inc.invalidate(0)
    assert inc(0) == 1
    assert calls == [0, 1, 0]

def test_cache_timedout():
    calls = []

    @cache(timeout=0)
    def inc(x):
        calls.append(x)
        return x + 1

    assert inc(0) == 1
    assert inc(0) == 1
    assert calls == [0, 0]
<commit_msg>Test @cache with mixed positional and keywords args<commit_after>from math import sin, cos
import pytest

from funcy.calc import *

def test_memoize():
    @memoize
    def inc(x):
        calls.append(x)
        return x + 1

    calls = []
    assert inc(0) == 1
    assert inc(1) == 2
    assert inc(0) == 1
    assert calls == [0, 1]

    # using kwargs
    assert inc(x=0) == 1
    assert inc(x=1) == 2
    assert inc(x=0) == 1
    assert calls == [0, 1, 0, 1]

def test_memoize_args_kwargs():
    @memoize
    def mul(x, by=1):
        calls.append((x, by))
        return x * by

    calls = []
    assert mul(0) == 0
    assert mul(1) == 1
    assert mul(0) == 0
    assert calls == [(0, 1), (1, 1)]

    # more with kwargs
    assert mul(0, 1) == 0
    assert mul(1, 1) == 1
    assert mul(0, 1) == 0
    assert calls == [(0, 1), (1, 1), (0, 1), (1, 1)]

def test_memoize_memory():
    @memoize
    def inc(x):
        calls.append(x)
        return x + 1

    calls = []
    inc(0)
    inc.memory.clear()
    inc(0)
    assert calls == [0, 0]

def test_memoize_key_func():
    @memoize(key_func=len)
    def inc(s):
        calls.append(s)
        return s * 2

    calls = []
    assert inc('a') == 'aa'
    assert inc('b') == 'aa'
    inc('ab')
    assert calls == ['a', 'ab']

def test_make_lookuper():
    @make_lookuper
    def letter_index():
        return ((c, i) for i, c in enumerate('abcdefghij'))

    assert letter_index('c') == 2
    with pytest.raises(LookupError): letter_index('_')

def test_make_lookuper_nested():
    tables_built = [0]

    @make_lookuper
    def function_table(f):
        tables_built[0] += 1
        return ((x, f(x)) for x in range(10))

    assert function_table(sin)(5) == sin(5)
    assert function_table(cos)(3) == cos(3)
    assert function_table(sin)(3) == sin(3)
    assert tables_built[0] == 2

    with pytest.raises(LookupError): function_table(cos)(-1)

def test_silent_lookuper():
    @silent_lookuper
    def letter_index():
        return ((c, i) for i, c in enumerate('abcdefghij'))

    assert letter_index('c') == 2
    assert letter_index('_') is None

def test_silnent_lookuper_nested():
    @silent_lookuper
    def function_table(f):
        return ((x, f(x)) for x in range(10))

    assert function_table(sin)(5) == sin(5)
    assert function_table(cos)(-1) is None

def test_cache():
    calls = []

    @cache(timeout=60)
    def inc(x):
        calls.append(x)
        return x + 1

    assert inc(0) == 1
    assert inc(1) == 2
    assert inc(0) == 1
    inc.invalidate(0)
    assert inc(0) == 1
    assert calls == [0, 1, 0]

def test_cache_mixed_args():
    @cache(timeout=60)
    def add(x, y):
        return x + y

    assert add(1, y=2) == 3

def test_cache_timedout():
    calls = []

    @cache(timeout=0)
    def inc(x):
        calls.append(x)
        return x + 1

    assert inc(0) == 1
    assert inc(0) == 1
    assert calls == [0, 0]
<|endoftext|>