<commit_before>"""Governs crud operations for profiles in the identity vault.
user_profile must be passed to this in the form required by dynamodb
{
    'id': 'ad|foo',
    'uuid': 'dc8cbad4-4426-406f-998e-9d95edb06bdc',
    'primary_email': 'foo@xyz.com',
    'sequence_number': '123456',
    'primary_email': 'foomcbar',
    'profile': 'jsondumpofuserfullprofile'
}
"""
import json
import logging
import uuid
from boto3.dynamodb.conditions import Attr
from boto3.dynamodb.conditions import Key
from botocore.exceptions import ClientError
from botocore.exceptions import ParamValidationError
from traceback import format_exc

logger = logging.getLogger(__name__)

class Profile(object):
    def __init__(self, dynamodb_table_resource=None, dynamodb_client=None, transactions=True):
        """Take a dynamodb table resource to use for operations."""
        self.table = dynamodb_table_resource
        self.client = dynamodb_client
        self.transactions = transactions

    def _run_transaction(self, transact_items):
        sequence_numbers = []
        for t in transact_items:
            if "Update" in t:
                sequence_numbers.append(t["Update"]["ExpressionAttributeValues"][":sn"]["S"])
            else:
                sequence_numbers.append(t["Put"]["Item"]["sequence_number"]["S"])

        try:
            self.client.transact_write_items(
                TransactItems=transact_items, ReturnConsumedCapacity="TOTAL", ReturnItemCollectionMetrics="SIZE"
            )
        except ClientError as e:
            logger.warning("Transaction failed", extra={"reason": e, "trace": format_exc()})
            raise ValueError("Transaction failed - profile issue?", e)

        return {"status": "200", "ResponseMetadata": {"HTTPStatusCode": 200}, "sequence_numbers": sequence_numbers}

    def create(self, user_profile):
        if self.transactions:
            res = self._create_with_transaction(user_profile)
        else:
            res = self._create_without_transaction(user_profile)

        if res.get("ResponseMetadata", False):
            status_code = res["ResponseMetadata"]["HTTPStatusCode"]
            sequence_number = user_profile["sequence_number"]

        return {"status": status_code, "sequence_number": sequence_number}

    def _create_without_transaction(self, user_profile):
        if user_profile["sequence_number"] is None:
            user_profile["sequence_number"] = str(uuid.uuid4().int)
        return self.table.put_item(
            Item={
                "id": user_profile["id"],
                "user_uuid": user_profile["user_uuid"],
                "profile": user_profile["profile"],
                "primary_email": user_profile["primary_email"],
                "primary_username": user_profile["primary_username"],
                "sequence_number": user_profile["sequence_number"],
                "active": bool(json.loads(user_profile["profile"])["active"]["value"])
            }
        )

    def _create_with_transaction(self, user_profile):
        if user_profile["sequence_number"] is None:
            user_profile["sequence_number"] = str(uuid.uuid4().int)
        transact_items = {
            "Put": {
                "Item": {
                    "id": {"S": user_profile["id"]},
                    "user_uuid": {"S": user_profile["user_uuid"]},
                    "profile": {"S": user_profile["profile"]},
                    "primary_email": {"S": user_profile["primary_email"]},
                    "primary_username": {"S": user_profile["primary_username"]},
                    "sequence_number": {"S": user_profile["sequence_number"]},
                    "active": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                },
                "ConditionExpression": "attribute_not_exists(id)",
                "TableName": self.table.name,
                "ReturnValuesOnConditionCheckFailure": "NONE",
            }
        }
        return self._run_transaction([transact_items])

    def update(self, user_profile):
        if self.transactions:
            res = self._update_with_transaction(user_profile)
        else:
            res = self._update_without_transaction(user_profile)
        return res

    def _update_with_transaction(self, user_profile):
        transact_items = {
            "Update": {
                "Key": {"id": {"S": user_profile["id"]}},
                "ExpressionAttributeValues": {
                    ":p": {"S": user_profile["profile"]},
                    ":u": {"S": user_profile["user_uuid"]},
                    ":pe": {"S": user_profile["primary_email"]},
                    ":pn": {"S": user_profile["primary_username"]},
                    ":sn": {"S": user_profile["sequence_number"]},
                    ":a": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                },
                "ConditionExpression": "attribute_exists(id)",
                "UpdateExpression": "SET profile = :p, primary_email = :pe, sequence_number = :sn, user_uuid = :u, primary_username = :pn, active = :a",
                "TableName": self.table.name,
                "ReturnValuesOnConditionCheckFailure": "NONE",
            }
        }
        return self._run_transaction([transact_items])

    def _update_without_transaction(self, user_profile):
        return self.table.put_item(
            Item={
                "id": user_profile["id"],
                "user_uuid": user_profile["user_uuid"],
                "profile": user_profile["profile"],
                "primary_email": user_profile["primary_email"],
                "primary_username": user_profile["primary_username"],
                "sequence_number": user_profile["sequence_number"],
                "active": bool(json.loads(user_profile["profile"])["active"]["value"])
            }
        )

    def delete(self, user_profile):
        res = self._delete_without_transaction(user_profile)
        return res

    def _delete_without_transaction(self, user_profile):
        return self.table.delete_item(Key={"id": user_profile["id"]})

    def create_batch(self, list_of_profiles):
        sequence_numbers = []
        for profile in list_of_profiles:
            sequence_numbers.append(profile["sequence_number"])
        if self.transactions:
            res = self._create_items_with_transaction(list_of_profiles)

            if res.get("ResponseMetadata", False):
                status_code = res["ResponseMetadata"]["HTTPStatusCode"]

            return {"status": status_code, "sequence_numbers": sequence_numbers}
        else:
            try:
                res = self._put_items_without_transaction(list_of_profiles)
                return {"status": "200", "sequence_numbers": sequence_numbers}
            except Exception as e:
                logger.error("Could not write batch due to: {}".format(e))
                return {"status": "500", "sequence_numbers": sequence_numbers}

    def _put_items_without_transaction(self, list_of_profiles):
        sequence_numbers = []
        with self.table.batch_writer() as batch:
            for profile in list_of_profiles:
                batch.put_item(
                    Item={
                        "id": user_profile["id"],
                        "user_uuid": user_profile["user_uuid"],
                        "profile": user_profile["profile"],
                        "primary_email": user_profile["primary_email"],
                        "primary_username": user_profile["primary_username"],
                        "sequence_number": user_profile["sequence_number"],
                        "active": bool(json.loads(user_profile["profile"])["active"]["value"])
                    }
                )
                sequence_numbers.append(profile["sequence_number"])

        return {"status": "200", "ResponseMetadata": {"HTTPStatusCode": 200}, "sequence_numbers": sequence_numbers}

    def _create_items_with_transaction(self, list_of_profiles):
        transact_items = []
        for user_profile in list_of_profiles:
            if user_profile["sequence_number"] is None:
                user_profile["sequence_number"] = str(uuid.uuid4().int)
            transact_item = {
                "Put": {
                    "Item": {
                        "id": {"S": user_profile["id"]},
                        "user_uuid": {"S": user_profile["user_uuid"]},
                        "profile": {"S": user_profile["profile"]},
                        "primary_email": {"S": user_profile["primary_email"]},
                        "primary_username": {"S": user_profile["primary_username"]},
                        "sequence_number": {"S": user_profile["sequence_number"]},
                        "active": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                    },
                    "ConditionExpression": "attribute_not_exists(id)",
                    "TableName": self.table.name,
                    "ReturnValuesOnConditionCheckFailure": "NONE",
                }
            }
            transact_items.append(transact_item)
        logger.debug("Attempting to create batch of transactions for: {}".format(transact_items))
        return self._run_transaction(transact_items)

    def update_batch(self, list_of_profiles):
        if self.transactions:
            res = self._update_batch_with_transaction(list_of_profiles)
        else:
            res = self._put_items_without_transaction(list_of_profiles)
        return res

    def _update_batch_with_transaction(self, list_of_profiles):
        transact_items = []
        for user_profile in list_of_profiles:
            transact_item = {
                "Update": {
                    "Key": {"id": {"S": user_profile["id"]}},
                    "ExpressionAttributeValues": {
                        ":p": {"S": user_profile["profile"]},
                        ":u": {"S": user_profile["user_uuid"]},
                        ":pe": {"S": user_profile["primary_email"]},
                        ":pn": {"S": user_profile["primary_username"]},
                        ":sn": {"S": user_profile["sequence_number"]},
                        ":a": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                    },
                    "ConditionExpression": "attribute_exists(id)",
                    "UpdateExpression": "SET profile = :p, primary_email = :pe, sequence_number = :sn, user_uuid = :u, primary_username = :pn, active = :a",
                    "TableName": self.table.name,
                    "ReturnValuesOnConditionCheckFailure": "NONE",
                }
            }
            transact_items.append(transact_item)
        logger.debug("Attempting to update batch of transactions for: {}".format(transact_items))
        return self._run_transaction(transact_items)

    def find_by_id(self, id):
        result = self.table.query(KeyConditionExpression=Key("id").eq(id))
        return result

    def find_by_email(self, primary_email):
        result = self.table.query(
            IndexName="{}-primary_email".format(self.table.table_name),
            KeyConditionExpression=Key("primary_email").eq(primary_email),
        )
        return result

    def find_by_uuid(self, uuid):
        result = self.table.query(
            IndexName="{}-user_uuid".format(self.table.table_name), KeyConditionExpression=Key("user_uuid").eq(uuid)
        )
        return result

    def find_by_username(self, primary_username):
        result = self.table.query(
            IndexName="{}-primary_username".format(self.table.table_name),
            KeyConditionExpression=Key("primary_username").eq(primary_username),
        )
        return result

    @property
    def all(self):
        response = self.table.scan()
        users = response.get("Items")
        while "LastEvaluatedKey" in response:
            response = self.table.scan(ExclusiveStartKey=response["LastEvaluatedKey"])
            users.extend(response["Items"])
        return users

    def _get_segment(self, segment=0, total_segments=10, connection_method=None, active=None):
        logger.info('Getting segment: {} of total: {}'.format(segment, total_segments))

        if connection_method:
            expression_attr = {
                ":id": {"S": connection_method}
            }
            filter_expression = "begins_with(id, :id)"

        if connection_method and active is not None:
            logger.info("Asking for only the users with active state: {}".format(active))
            expression_attr = {
                ":id": {"S": connection_method},
                ":a" : {"BOOL": active},
            }
            filter_expression = ":a = active AND begins_with(id, :id) AND attribute_exists(active)"

        response = self.client.scan(
            TableName=self.table.name,
            TotalSegments=total_segments,
            Segment=segment,
            ProjectionExpression="id, primary_email, user_uuid, active",
            FilterExpression=filter_expression,
            ExpressionAttributeValues=expression_attr,
        )
        users = response.get("Items", [])
        while "LastEvaluatedKey" in response:
            response = self.client.scan(
                TableName=self.table.name,
                TotalSegments=total_segments,
                Segment=segment,
                ProjectionExpression="id, primary_email, user_uuid, active",
                FilterExpression=filter_expression,
                ExpressionAttributeValues=expression_attr,
                ExclusiveStartKey=response["LastEvaluatedKey"],
            )
            users.extend(response["Items"])
        return users

    def all_filtered(self, connection_method=None, active=None):
        """
        @query_filter str login_method
        Returns a dict of all users filtered by query_filter
        """
        pool = []
        # We're choosing to divide the table in 3, then...
        pool_size = 5
        users = []
        for x in range(pool_size):
            # XXX TBD async this with asyncio
            users.extend(self._get_segment(segment=x, total_segments=pool_size, connection_method=connection_method, active=active))
        return users

    def find_or_create(self, user_profile):
        profilev2 = json.loads(user_profile["profile"])
        if len(self.find_by_id(profilev2["user_id"]["value"])["Items"]) > 0:
            res = self.update(user_profile)
            logger.info("A user profile exists already for: {}".format(profilev2["user_id"]["value"]))
        else:
            res = self.create(user_profile)
            logger.info("A user profile does not exist for: {}".format(profilev2["user_id"]["value"]))
        return res

    def find_or_create_batch(self, user_profiles):
        updates = []
        creations = []
        for user_profile in user_profiles:
            profilev2 = json.loads(user_profile["profile"])
            if len(self.find_by_id(profilev2["user_id"]["value"])["Items"]) > 0:
                logger.debug("Adding profile to the list of updates to perform: {}".format(profilev2))
                updates.append(user_profile)
            else:
                logger.debug("Adding profile to the list of creations to perform: {}".format(profilev2))
                creations.append(user_profile)

        try:
            if len(creations) > 0:
                res_create = self.create_batch(creations)
                logger.info("There are {} creations to perform in this batch.".format(res_create))
            else:
                res_create = None
        except ClientError as e:
            res_create = None
            logger.error("Could not run batch transaction due to: {}".format(e))
        except ParamValidationError as e:
            res_create = None
            logger.error("Could not run batch transaction due to: {}".format(e))

        try:
            if len(updates) > 0:
                res_update = self.update_batch(updates)
                logger.debug("There are {} updates to perform in this batch.".format(len(updates)))
            else:
                res_update = None
        except ClientError as e:
            res_update = None
            logger.error("Could not run batch transaction due to: {}".format(e))
        except ParamValidationError as e:
            res_update = None
            logger.error("Could not run batch transaction due to: {}".format(e))

        logger.info("Updates were: {}".format(len(updates)))
        logger.info("Creates were: {}".format(len(creations)))

        return [res_create, res_update]

    def all_by_page(self, next_page=None, limit=25):
        if next_page is not None:
            response = self.table.scan(Limit=limit, ExclusiveStartKey=next_page)
        else:
            response = self.table.scan(Limit=limit)
        return response

<commit_msg>fix bug introduced in batch write without transact<commit_after>"""Governs crud operations for profiles in the identity vault.
user_profile must be passed to this in the form required by dynamodb
{
    'id': 'ad|foo',
    'uuid': 'dc8cbad4-4426-406f-998e-9d95edb06bdc',
    'primary_email': 'foo@xyz.com',
    'sequence_number': '123456',
    'primary_email': 'foomcbar',
    'profile': 'jsondumpofuserfullprofile'
}
"""
import json
import logging
import uuid
from boto3.dynamodb.conditions import Attr
from boto3.dynamodb.conditions import Key
from botocore.exceptions import ClientError
from botocore.exceptions import ParamValidationError
from traceback import format_exc

logger = logging.getLogger(__name__)

class Profile(object):
    def __init__(self, dynamodb_table_resource=None, dynamodb_client=None, transactions=True):
        """Take a dynamodb table resource to use for operations."""
        self.table = dynamodb_table_resource
        self.client = dynamodb_client
        self.transactions = transactions

    def _run_transaction(self, transact_items):
        sequence_numbers = []
        for t in transact_items:
            if "Update" in t:
                sequence_numbers.append(t["Update"]["ExpressionAttributeValues"][":sn"]["S"])
            else:
                sequence_numbers.append(t["Put"]["Item"]["sequence_number"]["S"])

        try:
            self.client.transact_write_items(
                TransactItems=transact_items, ReturnConsumedCapacity="TOTAL", ReturnItemCollectionMetrics="SIZE"
            )
        except ClientError as e:
            logger.warning("Transaction failed", extra={"reason": e, "trace": format_exc()})
            raise ValueError("Transaction failed - profile issue?", e)

        return {"status": "200", "ResponseMetadata": {"HTTPStatusCode": 200}, "sequence_numbers": sequence_numbers}

    def create(self, user_profile):
        if self.transactions:
            res = self._create_with_transaction(user_profile)
        else:
            res = self._create_without_transaction(user_profile)

        if res.get("ResponseMetadata", False):
            status_code = res["ResponseMetadata"]["HTTPStatusCode"]
            sequence_number = user_profile["sequence_number"]

        return {"status": status_code, "sequence_number": sequence_number}

    def _create_without_transaction(self, user_profile):
        if user_profile["sequence_number"] is None:
            user_profile["sequence_number"] = str(uuid.uuid4().int)
        return self.table.put_item(
            Item={
                "id": user_profile["id"],
                "user_uuid": user_profile["user_uuid"],
                "profile": user_profile["profile"],
                "primary_email": user_profile["primary_email"],
                "primary_username": user_profile["primary_username"],
                "sequence_number": user_profile["sequence_number"],
                "active": bool(json.loads(user_profile["profile"])["active"]["value"])
            }
        )

    def _create_with_transaction(self, user_profile):
        if user_profile["sequence_number"] is None:
            user_profile["sequence_number"] = str(uuid.uuid4().int)
        transact_items = {
            "Put": {
                "Item": {
                    "id": {"S": user_profile["id"]},
                    "user_uuid": {"S": user_profile["user_uuid"]},
                    "profile": {"S": user_profile["profile"]},
                    "primary_email": {"S": user_profile["primary_email"]},
                    "primary_username": {"S": user_profile["primary_username"]},
                    "sequence_number": {"S": user_profile["sequence_number"]},
                    "active": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                },
                "ConditionExpression": "attribute_not_exists(id)",
                "TableName": self.table.name,
                "ReturnValuesOnConditionCheckFailure": "NONE",
            }
        }
        return self._run_transaction([transact_items])

    def update(self, user_profile):
        if self.transactions:
            res = self._update_with_transaction(user_profile)
        else:
            res = self._update_without_transaction(user_profile)
        return res

    def _update_with_transaction(self, user_profile):
        transact_items = {
            "Update": {
                "Key": {"id": {"S": user_profile["id"]}},
                "ExpressionAttributeValues": {
                    ":p": {"S": user_profile["profile"]},
                    ":u": {"S": user_profile["user_uuid"]},
                    ":pe": {"S": user_profile["primary_email"]},
                    ":pn": {"S": user_profile["primary_username"]},
                    ":sn": {"S": user_profile["sequence_number"]},
                    ":a": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                },
                "ConditionExpression": "attribute_exists(id)",
                "UpdateExpression": "SET profile = :p, primary_email = :pe, sequence_number = :sn, user_uuid = :u, primary_username = :pn, active = :a",
                "TableName": self.table.name,
                "ReturnValuesOnConditionCheckFailure": "NONE",
            }
        }
        return self._run_transaction([transact_items])

    def _update_without_transaction(self, user_profile):
        return self.table.put_item(
            Item={
                "id": user_profile["id"],
                "user_uuid": user_profile["user_uuid"],
                "profile": user_profile["profile"],
                "primary_email": user_profile["primary_email"],
                "primary_username": user_profile["primary_username"],
                "sequence_number": user_profile["sequence_number"],
                "active": bool(json.loads(user_profile["profile"])["active"]["value"])
            }
        )

    def delete(self, user_profile):
        res = self._delete_without_transaction(user_profile)
        return res

    def _delete_without_transaction(self, user_profile):
        return self.table.delete_item(Key={"id": user_profile["id"]})

    def create_batch(self, list_of_profiles):
        sequence_numbers = []
        for profile in list_of_profiles:
            sequence_numbers.append(profile["sequence_number"])
        if self.transactions:
            res = self._create_items_with_transaction(list_of_profiles)

            if res.get("ResponseMetadata", False):
                status_code = res["ResponseMetadata"]["HTTPStatusCode"]

            return {"status": status_code, "sequence_numbers": sequence_numbers}
        else:
            try:
                res = self._put_items_without_transaction(list_of_profiles)
                return {"status": "200", "sequence_numbers": sequence_numbers}
            except Exception as e:
                logger.error("Could not write batch due to: {}".format(e))
                return {"status": "500", "sequence_numbers": sequence_numbers}

    def _put_items_without_transaction(self, list_of_profiles):
        sequence_numbers = []
        with self.table.batch_writer() as batch:
            for profile in list_of_profiles:
                batch.put_item(
                    Item={
                        "id": profile["id"],
                        "user_uuid": profile["user_uuid"],
                        "profile": profile["profile"],
                        "primary_email": profile["primary_email"],
                        "primary_username": profile["primary_username"],
                        "sequence_number": profile["sequence_number"],
                        "active": bool(json.loads(profile["profile"])["active"]["value"])
                    }
                )
                sequence_numbers.append(profile["sequence_number"])

        return {"status": "200", "ResponseMetadata": {"HTTPStatusCode": 200}, "sequence_numbers": sequence_numbers}

    def _create_items_with_transaction(self, list_of_profiles):
        transact_items = []
        for user_profile in list_of_profiles:
            if user_profile["sequence_number"] is None:
                user_profile["sequence_number"] = str(uuid.uuid4().int)
            transact_item = {
                "Put": {
                    "Item": {
                        "id": {"S": user_profile["id"]},
                        "user_uuid": {"S": user_profile["user_uuid"]},
                        "profile": {"S": user_profile["profile"]},
                        "primary_email": {"S": user_profile["primary_email"]},
                        "primary_username": {"S": user_profile["primary_username"]},
                        "sequence_number": {"S": user_profile["sequence_number"]},
                        "active": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                    },
                    "ConditionExpression": "attribute_not_exists(id)",
                    "TableName": self.table.name,
                    "ReturnValuesOnConditionCheckFailure": "NONE",
                }
            }
            transact_items.append(transact_item)
        logger.debug("Attempting to create batch of transactions for: {}".format(transact_items))
        return self._run_transaction(transact_items)

    def update_batch(self, list_of_profiles):
        if self.transactions:
            res = self._update_batch_with_transaction(list_of_profiles)
        else:
            res = self._put_items_without_transaction(list_of_profiles)
        return res

    def _update_batch_with_transaction(self, list_of_profiles):
        transact_items = []
        for user_profile in list_of_profiles:
            transact_item = {
                "Update": {
                    "Key": {"id": {"S": user_profile["id"]}},
                    "ExpressionAttributeValues": {
                        ":p": {"S": user_profile["profile"]},
                        ":u": {"S": user_profile["user_uuid"]},
                        ":pe": {"S": user_profile["primary_email"]},
                        ":pn": {"S": user_profile["primary_username"]},
                        ":sn": {"S": user_profile["sequence_number"]},
                        ":a": {"BOOL": json.loads(user_profile["profile"])["active"]["value"]},
                    },
                    "ConditionExpression": "attribute_exists(id)",
                    "UpdateExpression": "SET profile = :p, primary_email = :pe, sequence_number = :sn, user_uuid = :u, primary_username = :pn, active = :a",
                    "TableName": self.table.name,
                    "ReturnValuesOnConditionCheckFailure": "NONE",
                }
            }
            transact_items.append(transact_item)
        logger.debug("Attempting to update batch of transactions for: {}".format(transact_items))
        return self._run_transaction(transact_items)

    def find_by_id(self, id):
        result = self.table.query(KeyConditionExpression=Key("id").eq(id))
        return result

    def find_by_email(self, primary_email):
        result = self.table.query(
            IndexName="{}-primary_email".format(self.table.table_name),
            KeyConditionExpression=Key("primary_email").eq(primary_email),
        )
        return result

    def find_by_uuid(self, uuid):
        result = self.table.query(
            IndexName="{}-user_uuid".format(self.table.table_name), KeyConditionExpression=Key("user_uuid").eq(uuid)
        )
        return result

    def find_by_username(self, primary_username):
        result = self.table.query(
            IndexName="{}-primary_username".format(self.table.table_name),
            KeyConditionExpression=Key("primary_username").eq(primary_username),
        )
        return result

    @property
    def all(self):
        response = self.table.scan()
        users = response.get("Items")
        while "LastEvaluatedKey" in response:
            response = self.table.scan(ExclusiveStartKey=response["LastEvaluatedKey"])
            users.extend(response["Items"])
        return users

    def _get_segment(self, segment=0, total_segments=10, connection_method=None, active=None):
        logger.info('Getting segment: {} of total: {}'.format(segment, total_segments))

        if connection_method:
            expression_attr = {
                ":id": {"S": connection_method}
            }
            filter_expression = "begins_with(id, :id)"

        if connection_method and active is not None:
            logger.info("Asking for only the users with active state: {}".format(active))
            expression_attr = {
                ":id": {"S": connection_method},
                ":a" : {"BOOL": active},
            }
            filter_expression = ":a = active AND begins_with(id, :id) AND attribute_exists(active)"

        response = self.client.scan(
            TableName=self.table.name,
            TotalSegments=total_segments,
            Segment=segment,
            ProjectionExpression="id, primary_email, user_uuid, active",
            FilterExpression=filter_expression,
            ExpressionAttributeValues=expression_attr,
        )
        users = response.get("Items", [])
        while "LastEvaluatedKey" in response:
            response = self.client.scan(
                TableName=self.table.name,
                TotalSegments=total_segments,
                Segment=segment,
                ProjectionExpression="id, primary_email, user_uuid, active",
                FilterExpression=filter_expression,
                ExpressionAttributeValues=expression_attr,
                ExclusiveStartKey=response["LastEvaluatedKey"],
            )
            users.extend(response["Items"])
        return users

    def all_filtered(self, connection_method=None, active=None):
        """
        @query_filter str login_method
        Returns a dict of all users filtered by query_filter
        """
        pool = []
        # We're choosing to divide the table in 3, then...
        pool_size = 5
        users = []
        for x in range(pool_size):
            # XXX TBD async this with asyncio
            users.extend(self._get_segment(segment=x, total_segments=pool_size, connection_method=connection_method, active=active))
        return users

    def find_or_create(self, user_profile):
        profilev2 = json.loads(user_profile["profile"])
        if len(self.find_by_id(profilev2["user_id"]["value"])["Items"]) > 0:
            res = self.update(user_profile)
            logger.info("A user profile exists already for: {}".format(profilev2["user_id"]["value"]))
        else:
            res = self.create(user_profile)
            logger.info("A user profile does not exist for: {}".format(profilev2["user_id"]["value"]))
        return res

    def find_or_create_batch(self, user_profiles):
        updates = []
        creations = []
        for user_profile in user_profiles:
            profilev2 = json.loads(user_profile["profile"])
            if len(self.find_by_id(profilev2["user_id"]["value"])["Items"]) > 0:
                logger.debug("Adding profile to the list of updates to perform: {}".format(profilev2))
                updates.append(user_profile)
            else:
                logger.debug("Adding profile to the list of creations to perform: {}".format(profilev2))
                creations.append(user_profile)

        try:
            if len(creations) > 0:
                res_create = self.create_batch(creations)
                logger.info("There are {} creations to perform in this batch.".format(res_create))
            else:
                res_create = None
        except ClientError as e:
            res_create = None
            logger.error("Could not run batch transaction due to: {}".format(e))
        except ParamValidationError as e:
            res_create = None
            logger.error("Could not run batch transaction due to: {}".format(e))

        try:
            if len(updates) > 0:
                res_update = self.update_batch(updates)
                logger.debug("There are {} updates to perform in this batch.".format(len(updates)))
            else:
                res_update = None
        except ClientError as e:
            res_update = None
            logger.error("Could not run batch transaction due to: {}".format(e))
        except ParamValidationError as e:
            res_update = None
            logger.error("Could not run batch transaction due to: {}".format(e))

        logger.info("Updates were: {}".format(len(updates)))
        logger.info("Creates were: {}".format(len(creations)))

        return [res_create, res_update]

    def all_by_page(self, next_page=None, limit=25):
        if next_page is not None:
            response = self.table.scan(Limit=limit, ExclusiveStartKey=next_page)
        else:
            response = self.table.scan(Limit=limit)
        return response

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from mx import DateTime
from mx.DateTime import now
import time
import tools
import netsvc
from osv import fields, osv
import ir
from tools.translate import _

import sys
from tools.translate import _

try:
    from lxml import etree
except ImportError:
    sys.stderr.write("ERROR: Import lxml module\n")
    sys.stderr.write("ERROR: Try to install the python-lxml package\n")

class project_gtd_context(osv.osv):
    _name = "project.gtd.context"
    _description = "Context"
    _columns = {
        'name': fields.char('Context', size=64, required=True, select=1),
        'sequence': fields.integer('Sequence', help="Gives the sequence order when displaying a list of contexts."),
        'project_default_id': fields.many2one('project.project', 'Default Project', required=True),
    }
    _defaults = {
        'sequence': 1
    }
    _order = "sequence, name"

project_gtd_context()

class project_gtd_timebox(osv.osv):
    _name = "project.gtd.timebox"
    _order = "sequence"
    _columns = {
        'name': fields.char('Timebox', size=64, required=True, select=1),
        'sequence': fields.integer('Sequence', help="Gives the sequence order when displaying a list of timebox."),
        'icon': fields.selection(tools.icons, 'Icon', size=64),
    }

#    def fields_view_get(self, cr, uid, view_id=None, view_type='form', context=None, toolbar=False, submenu=False):
#        res = super(project_gtd_timebox,self).fields_view_get(cr, uid, view_id, view_type, context, toolbar=toolbar, submenu=submenu)
#        if (res['type']=='form') and ('record_id' in context):
#            if context['record_id']:
#                rec = self.browse(cr, uid, int(context['record_id']), context)
#            else:
#                iids = self.search(cr,uid, [('user_id','=',uid),('parent_id','=',False)], context=context)
#                if len(iids):
#                    rec = self.browse(cr, uid, int(iids[0]), context=context)
#                else:
#                    return res
#            res['arch'] = """
#    <form string="Daily Timebox">
#        <field name="name" readonly="1"/>
#        <notebook position="top">
#            """
#            for i in range(1,7):
#                if not getattr(rec, 'context%d_id'%i):
#                    continue
#                res['arch']+= """
#            <page string="%s">
#                <field name="%s" colspan="4" nolabel="1">
#                    <tree editable="bottom" colors="grey:state in ('done','pending');red:state=='cancelled'" string="Tasks">
#                        <field name="name"/>
#                """ % (getattr(rec, 'context%d_id'%(i,)).name.encode('utf-8'), 'task%d_ids'%(i,))
#                if rec.col_project:
#                    res['arch'] += '<field name="project_id" required="1"/>\n'
#                if rec.col_date_start:
#                    res['arch'] += '<field name="date_start"/>\n'
#                if rec.col_priority:
#                    res['arch'] += '<field name="priority"/>\n'
#                if rec.col_deadline:
#                    res['arch'] += '<field name="date_deadline"/>\n'
#                if rec.col_planned_hours:
#                    res['arch'] += '<field name="planned_hours"  widget="float_time" sum="Est. Hours"/>\n'
#                if rec.col_effective_hours:
#                    res['arch'] += '<field name="effective_hours"  widget="float_time" sum="%s"/>\n' % (_('Eff. Hours'),)
#                res['arch'] += """
#                        <field name="state" readonly="1"/>
#                    </tree>
#                </field>
#            </page>
#                """
#            res['arch']+="""
#        </notebook>
#    </form>
#            """
#        doc = etree.XML(res['arch'])
#        xarch, xfields = self._view_look_dom_arch(cr, uid, doc, view_id, context=context)
#        res['arch'] = xarch
#        res['fields'] = xfields
#        return res

project_gtd_timebox()

class project_task(osv.osv):
    _inherit = "project.task"
    _columns = {
        'timebox_id': fields.many2one('project.gtd.timebox', "Timebox",help="Time-laps during which task has to be treated"),
        'context_id': fields.many2one('project.gtd.context', "Context",help="The context place where user has to treat task"),
     }

    def copy_data(self, cr, uid, id, default=None, context=None):
        if context is None:
            context = {}
        if not default:
            default = {}
        default['timebox_id'] = False
        default['context_id'] = False
        return super(project_task,self).copy_data(cr, uid, id, default, context)

    def _get_context(self,cr, uid, context=None):
        ids = self.pool.get('project.gtd.context').search(cr, uid, [], context=context)
        return ids and ids[0] or False

    _defaults = {
        'context_id': _get_context
    }
    def next_timebox(self, cr, uid, ids, *args):
        timebox_obj = self.pool.get('project.gtd.timebox')
        timebox_ids = timebox_obj.search(cr,uid,[])
        if not timebox_ids: return True
        for task in self.browse(cr,uid,ids):
            timebox = task.timebox_id.id
            if not timebox:
                self.write(cr, uid, task.id, {'timebox_id': timebox_ids[0]})
            elif timebox_ids.index(timebox) != len(timebox_ids)-1:
                index = timebox_ids.index(timebox)
                self.write(cr, uid, task.id, {'timebox_id': timebox_ids[index+1]})
        return True

    def prev_timebox(self, cr, uid, ids, *args):
        timebox_obj = self.pool.get('project.gtd.timebox')
        timebox_ids = timebox_obj.search(cr,uid,[])
        for task in self.browse(cr,uid,ids):
            timebox = task.timebox_id.id
            if timebox:
                if timebox_ids.index(timebox):
                    index = timebox_ids.index(timebox)
                    self.write(cr, uid, task.id, {'timebox_id': timebox_ids[index - 1]})
                else:
                    self.write(cr, uid, task.id, {'timebox_id': False})
        return True

    def fields_view_get(self, cr, uid, view_id=None, view_type='form', context=None, toolbar=False, submenu=False):
        res = super(project_task,self).fields_view_get(cr, uid, view_id, view_type, context, toolbar=toolbar, submenu=submenu)
        search_extended = False
        timebox_obj = self.pool.get('project.gtd.timebox')
        if res['type'] == 'search':
            tt = timebox_obj.browse(cr, uid, timebox_obj.search(cr,uid,[]))
            search_extended ='''<newline/><group col="%d" expand="0" string="%s" groups="project_gtd.group_project_getting">''' % (len(tt)+7,_('Getting Things Done'))
            search_extended += '''<filter domain="[('timebox_id','=', False)]" context="{'set_editable':True,'set_visible':True}" icon="gtk-new" string="%s"/>''' % (_('Inbox'),)
            search_extended += '''<filter domain="[('state', 'in', ('draft','open','pending'))]" context="{'set_editable':True,'set_visible':True}" icon="gtk-new" string="%s"/>''' % (_('All'),)
            search_extended += '''<separator orientation="vertical"/>'''
            for time in tt:
                if time.icon:
                    icon = time.icon
                else :
                    icon=""
                search_extended += '''<filter domain="[('timebox_id','=', ''' + str(time.id) + ''')]" icon="''' + icon + '''" string="''' + time.name + '''" context="{'set_visible':True}"/>'''
            search_extended += '''
            <separator orientation="vertical"/>
            <field name="context_id" select="1" widget="selection" />
            <field name="priority" select="1" />
            </group>
            </search> '''
        if search_extended:
            res['arch'] = res['arch'].replace('</search>',search_extended)
            attrs_sel = self.pool.get('project.gtd.context').name_search(cr, uid, '', [], context=context)
            context_id_info = self.pool.get('project.task').fields_get(cr, uid, ['context_id'])
            context_id_info['context_id']['selection'] = attrs_sel
            res['fields'].update(context_id_info)
        return res

project_task()

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>imp<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from mx import DateTime
from mx.DateTime import now
import time
import tools
import netsvc
from osv import fields, osv
import ir
from tools.translate import _

import sys
from tools.translate import _

try:
    from lxml import etree
except ImportError:
    sys.stderr.write("ERROR: Import lxml module\n")
    sys.stderr.write("ERROR: Try to install the python-lxml package\n")

class project_gtd_context(osv.osv):
    _name = "project.gtd.context"
    _description = "Context"
    _columns = {
        'name': fields.char('Context', size=64, required=True, select=1),
        'sequence': fields.integer('Sequence', help="Gives the sequence order when displaying a list of contexts."),
        'project_default_id': fields.many2one('project.project', 'Default Project', required=True),
    }
    _defaults = {
        'sequence': 1
    }
    _order = "sequence, name"

project_gtd_context()

class project_gtd_timebox(osv.osv):
    _name = "project.gtd.timebox"
    _order = "sequence"
    _columns = {
        'name': fields.char('Timebox', size=64, required=True, select=1),
        'sequence': fields.integer('Sequence', help="Gives the sequence order when displaying a list of timebox."),
        'icon': fields.selection(tools.icons, 'Icon', size=64),
    }

#    def fields_view_get(self, cr, uid, view_id=None, view_type='form', context=None, toolbar=False, submenu=False):
#        res = super(project_gtd_timebox,self).fields_view_get(cr, uid, view_id, view_type, context, toolbar=toolbar, submenu=submenu)
#        if (res['type']=='form') and ('record_id' in context):
#            if context['record_id']:
#                rec = self.browse(cr, uid, int(context['record_id']), context)
#            else:
#                iids = self.search(cr,uid, [('user_id','=',uid),('parent_id','=',False)], context=context)
#                if len(iids):
#                    rec = self.browse(cr, uid, int(iids[0]), context=context)
#                else:
#                    return res
#            res['arch'] = """
#    <form string="Daily Timebox">
#        <field name="name" readonly="1"/>
#        <notebook position="top">
#            """
#            for i in range(1,7):
#                if not getattr(rec, 'context%d_id'%i):
#                    continue
#                res['arch']+= """
#            <page string="%s">
#                <field name="%s" colspan="4" nolabel="1">
#                    <tree editable="bottom" colors="grey:state in ('done','pending');red:state=='cancelled'" string="Tasks">
#                        <field name="name"/>
#                """ % (getattr(rec, 'context%d_id'%(i,)).name.encode('utf-8'), 'task%d_ids'%(i,))
#                if rec.col_project:
#                    res['arch'] += '<field name="project_id" required="1"/>\n'
#                if rec.col_date_start:
#                    res['arch'] += '<field name="date_start"/>\n'
#                if rec.col_priority:
#                    res['arch'] += '<field name="priority"/>\n'
#                if rec.col_deadline:
#                    res['arch'] += '<field name="date_deadline"/>\n'
#                if rec.col_planned_hours:
#                    res['arch'] += '<field name="planned_hours"  widget="float_time" sum="Est. Hours"/>\n'
#                if rec.col_effective_hours:
#                    res['arch'] += '<field name="effective_hours"  widget="float_time" sum="%s"/>\n' % (_('Eff. Hours'),)
#                res['arch'] += """
#                        <field name="state" readonly="1"/>
#                    </tree>
#                </field>
#            </page>
#                """
#            res['arch']+="""
#        </notebook>
#    </form>
#            """
#        doc = etree.XML(res['arch'])
#        xarch, xfields = self._view_look_dom_arch(cr, uid, doc, view_id, context=context)
#        res['arch'] = xarch
#        res['fields'] = xfields
#        return res

project_gtd_timebox()

class project_task(osv.osv):
    _inherit = "project.task"
    _columns = {
        'timebox_id': fields.many2one('project.gtd.timebox', "Timebox",help="Time-laps during which task has to be treated"),
        'context_id': fields.many2one('project.gtd.context', "Context",help="The context place where user has to treat task"),
     }

    def copy_data(self, cr, uid, id, default=None, context=None):
        if context is None:
            context = {}
        if not default:
            default = {}
        default['timebox_id'] = False
        default['context_id'] = False
        return super(project_task,self).copy_data(cr, uid, id, default, context)

    def _get_context(self,cr, uid, context=None):
        ids = self.pool.get('project.gtd.context').search(cr, uid, [], context=context)
        return ids and ids[0] or False

    _defaults = {
        'context_id': _get_context
    }
    def next_timebox(self, cr, uid, ids, *args):
        timebox_obj = self.pool.get('project.gtd.timebox')
        timebox_ids = timebox_obj.search(cr,uid,[])
        if not timebox_ids: return True
        for task in self.browse(cr,uid,ids):
            timebox = task.timebox_id.id
            if not timebox:
                self.write(cr, uid, task.id, {'timebox_id': timebox_ids[0]})
            elif timebox_ids.index(timebox) != len(timebox_ids)-1:
                index = timebox_ids.index(timebox)
                self.write(cr, uid, task.id, {'timebox_id': timebox_ids[index+1]})
        return True

    def prev_timebox(self, cr, uid, ids, *args):
        timebox_obj = self.pool.get('project.gtd.timebox')
        timebox_ids = timebox_obj.search(cr,uid,[])
        for task in self.browse(cr,uid,ids):
            timebox = task.timebox_id.id
            if timebox:
                if timebox_ids.index(timebox):
                    index = timebox_ids.index(timebox)
                    self.write(cr, uid, task.id, {'timebox_id': timebox_ids[index - 1]})
                else:
                    self.write(cr, uid, task.id, {'timebox_id': False})
        return True

    def fields_view_get(self, cr, uid, view_id=None, view_type='form', context=None, toolbar=False, submenu=False):
        res = super(project_task,self).fields_view_get(cr, uid, view_id, view_type, context, toolbar=toolbar, submenu=submenu)
        search_extended = False
        timebox_obj = self.pool.get('project.gtd.timebox')
        if res['type'] == 'search':
            tt = timebox_obj.browse(cr, uid, timebox_obj.search(cr,uid,[]))
            search_extended ='''<newline/><group col="%d" expand="0" string="%s" groups="project_gtd.group_project_getting">''' % (len(tt)+7,_('Getting Things Done'))
            search_extended += '''<filter domain="[('timebox_id','=', False)]" context="{'set_editable':True,'set_visible':True}" icon="gtk-new" string="%s"/>''' % (_('Inbox'),)
            search_extended += '''<filter domain="[('state', 'in', ('draft','open','pending'))]" context="{'set_editable':True,'set_visible':True}" icon="gtk-new" string="%s"/>''' % (_('Editable'),)
            search_extended += '''<separator orientation="vertical"/>'''
            for time in tt:
                if time.icon:
                    icon = time.icon
                else :
                    icon=""
                search_extended += '''<filter domain="[('timebox_id','=', ''' + str(time.id) + ''')]" icon="''' + icon + '''" string="''' + time.name + '''" context="{'set_visible':True}"/>'''
            search_extended += '''
            <separator orientation="vertical"/>
            <field name="context_id" select="1" widget="selection" />
            <field name="priority" select="1" />
            </group>
            </search> '''
        if search_extended:
            res['arch'] = res['arch'].replace('</search>',search_extended)
            attrs_sel = self.pool.get('project.gtd.context').name_search(cr, uid, '', [], context=context)
            context_id_info = self.pool.get('project.task').fields_get(cr, uid, ['context_id'])
            context_id_info['context_id']['selection'] = attrs_sel
            res['fields'].update(context_id_info)
        return res

project_task()

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<|endoftext|>


<commit_before>#
# Created as part of the StratusLab project (http://stratuslab.eu),
# co-funded by the European Commission under the Grant Agreement
# INFSO-RI-261552."
#
# Copyright (c) 2011, SixSq Sarl
#
# 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.
#
import os
import shutil
import tempfile
import urllib2
import hashlib
from stratuslab.ManifestInfo import _normalizeForInstanceAttribute as normalizeManifestElementForClassAttr

from stratuslab import Util
from stratuslab import Defaults
from stratuslab.Exceptions import ExecutionException
from stratuslab.ManifestInfo import ManifestInfo
from stratuslab.Signator import Signator
from stratuslab.Compressor import Compressor
from stratuslab.ConfigHolder import ConfigHolder
from stratuslab.Exceptions import InputException
from stratuslab.Exceptions import ValidationException

from Util import Util as MarketplaceUtil

etree = Util.importETree()

class Downloader(object):

    ENDPOINT = Defaults.marketplaceEndpoint
    LOCAL_IMAGE_FILENAME = '/tmp/image.img'

    def __init__(self, configHolder=ConfigHolder()):
        self.localImageFilename = ''
        self.configHolder = configHolder
        self.imageUrl = ''
        self.compression = ''
        self.marketplaceEndpoint = Downloader.ENDPOINT

        configHolder.assign(self)

        self.localImageFilename = os.path.abspath(self.localImageFilename)
        self.manifestObject = None

    def _getManifest(self, imageId, tempMetadataFilename):
        """Return manifest as ManifestInfo object.
        """
        url = MarketplaceUtil.metadataUrl(self.marketplaceEndpoint, imageId)
        try:
            return self.__getManifest(url, tempMetadataFilename)
        except:
            print "Failed to get manifest for image: %s" % imageId
            raise

    def __getManifest(self, url, tempMetadataFilename):
        """Return manifest as ManifestInfo object.
        """
        errorMessage = 'Failed to find metadata entry: %s' % url
        try:
            self._download(url, tempMetadataFilename)
        except urllib2.HTTPError:
            raise InputException(errorMessage)

        if not os.path.exists(tempMetadataFilename):
            raise InputException(errorMessage)

        manifestInfo = ManifestInfo(self.configHolder)

        try:
            manifestInfo.parseManifestFromFile(tempMetadataFilename)
        except ExecutionException, ex:
            raise InputException('Error parsing the metadata corresponding to url %s, with detail %s' % (url, ex))
        return manifestInfo

    def getImageLocations(self, imageId=''):
        return self.getImageElementValue('locations', imageId)

    def getImageVersion(self, imageId=''):
        return self.getImageElementValue('version', imageId)

    def getImageElementValue(self, element, imageId=''):
        if imageId:
            self._checkManifestAndImageId(imageId)
        elementNorm = normalizeManifestElementForClassAttr(element)
        try:
            return getattr(self.manifestObject, elementNorm)
        except AttributeError, ex:
            raise ExecutionException("Couldn't get '%s' element (normalized to '%s') from manifest. %s" %
                                     (element, elementNorm, str(ex)))

    def _checkManifestAndImageId(self, imageId):
        if not self.manifestObject:
            self.downloadManifestByImageId(imageId)
        if imageId != self.manifestObject.identifier:
            raise InputException('Given image ID [%s] does not match to downloaded [%s]' %
                                    (imageId, self.manifestObject.identifier))

    def downloadManifestByImageId(self, imageId):
        tempMetadataFilename = tempfile.mktemp()
        try:
            self.manifestObject = self._getManifest(imageId, tempMetadataFilename)
        finally:
            try:
                os.unlink(tempMetadataFilename)
            except:
                pass

    def download(self, uri):
        tempMetadataFilename = tempfile.mktemp()

        manifestInfo = self._getManifest(uri, tempMetadataFilename)

        tempImageFilename = ''
        for location in manifestInfo.locations:
            self._printDetail('Looking for image: %s' % location)
            try:
                tempImageFilename = self._downloadImage(location)
                break
            except KeyboardInterrupt:
                raise
            except:
                pass

        if not os.path.exists(tempImageFilename):
            raise InputException('Failed to find image matching metadata: %s' % uri)

        self._verifySignature(tempImageFilename, tempMetadataFilename)

        tempImageFilename = self._inflateImage(tempImageFilename)

        self._verifyHash(tempImageFilename, manifestInfo.sha1)

        shutil.copy2(tempImageFilename, self.localImageFilename)

        os.remove(tempImageFilename)
        os.remove(tempMetadataFilename)

        return self.localImageFilename

    def _loadDom(self, filename):
        file = open(filename).read()
        dom = etree.fromstring(file)
        return dom

    def _downloadImage(self, url):
        compressionExtension = self._extractCompressionExtension(url)

        localFilename = tempfile.mktemp()
        localImageName = localFilename + compressionExtension

        Util.wget(url, localImageName)

        return localImageName

    def _extractCompressionExtension(self, url):
        compression = url.split('.')[-1]

        if compression in Compressor.compressionFormats:
            return '.' + compression
        else:
            return ''

    def _download(self, url, localFilename):
        try:
            return Util.wget(url, localFilename)
        except urllib2.URLError, ex:
            raise InputException('Failed to download: %s, with detail: %s' % (url, str(ex)))

    def _verifySignature(self, imageFilename, metadataFilename):
        signator = Signator(metadataFilename, self.configHolder)
        res = signator.validate()
        if res:
            raise ExecutionException('Failed to validate metadata file')

    def _inflateImage(self, imageFilename):
        extension = self._extractCompressionExtension(imageFilename)
        inflatedFilename = imageFilename
        if extension:
            self._printDetail('Inflating image %s' % imageFilename)
            Compressor.inflate(imageFilename)
            inflatedFilename = imageFilename[:-len(extension)]

        return inflatedFilename

    def _verifyHash(self, imageFilename, hashFromManifest):

        data = open(imageFilename).read()
        sha1 = hashlib.sha1()
        sha1.update(data)

        imageHash = sha1.hexdigest()

        if imageHash != hashFromManifest:
            raise ValidationException('Manifest hash code different from downloaded image: image=%s, matadata=%s' % (imageHash, hashFromManifest))

    def _printDetail(self, message):
        Util.printDetail(message, self.verboseLevel, 1)

<commit_msg>Refactored imageId to resourceUri to make it clear that we're dealing with fully qualified resource uri, not image unique identifier, which can correspond to several manifest entries <commit_after>#
# Created as part of the StratusLab project (http://stratuslab.eu),
# co-funded by the European Commission under the Grant Agreement
# INFSO-RI-261552."
#
# Copyright (c) 2011, SixSq Sarl
#
# 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.
#
import os
import shutil
import tempfile
import urllib2
import hashlib
from stratuslab.ManifestInfo import _normalizeForInstanceAttribute as normalizeManifestElementForClassAttr

from stratuslab import Util
from stratuslab import Defaults
from stratuslab.Exceptions import ExecutionException
from stratuslab.ManifestInfo import ManifestInfo
from stratuslab.Signator import Signator
from stratuslab.Compressor import Compressor
from stratuslab.ConfigHolder import ConfigHolder
from stratuslab.Exceptions import InputException
from stratuslab.Exceptions import ValidationException

from Util import Util as MarketplaceUtil

etree = Util.importETree()

class Downloader(object):

    ENDPOINT = Defaults.marketplaceEndpoint
    LOCAL_IMAGE_FILENAME = '/tmp/image.img'

    def __init__(self, configHolder=ConfigHolder()):
        self.localImageFilename = ''
        self.configHolder = configHolder
        self.imageUrl = ''
        self.compression = ''
        self.marketplaceEndpoint = Downloader.ENDPOINT

        configHolder.assign(self)

        self.localImageFilename = os.path.abspath(self.localImageFilename)
        self.manifestObject = None

    def _getManifest(self, resourceUri, tempMetadataFilename):
        """Return manifest as ManifestInfo object.
        """
        url = MarketplaceUtil.metadataUrl(self.marketplaceEndpoint, resourceUri)
        try:
            return self.__getManifest(url, tempMetadataFilename)
        except:
            print "Failed to get manifest for resource uri: %s" % resourceUri
            raise

    def __getManifest(self, url, tempMetadataFilename):
        """Return manifest as ManifestInfo object.
        """
        errorMessage = 'Failed to find metadata entry: %s' % url
        try:
            self._download(url, tempMetadataFilename)
        except urllib2.HTTPError:
            raise InputException(errorMessage)

        if not os.path.exists(tempMetadataFilename):
            raise InputException(errorMessage)

        manifestInfo = ManifestInfo(self.configHolder)

        try:
            manifestInfo.parseManifestFromFile(tempMetadataFilename)
        except ExecutionException, ex:
            raise InputException('Error parsing the metadata corresponding to url %s, with detail %s' % (url, ex))
        return manifestInfo

    def getImageLocations(self, resourceUri=''):
        return self.getImageElementValue('locations', resourceUri)

    def getImageVersion(self, resourceUri=''):
        return self.getImageElementValue('version', resourceUri)

    def getImageElementValue(self, element, resourceUri=''):
        if resourceUri:
            self._checkManifestAndImageId(resourceUri)
        elementNorm = normalizeManifestElementForClassAttr(element)
        try:
            return getattr(self.manifestObject, elementNorm)
        except AttributeError, ex:
            raise ExecutionException("Couldn't get '%s' element (normalized to '%s') from manifest. %s" %
                                     (element, elementNorm, str(ex)))

    def _checkManifestAndImageId(self, resourceUri):
        if not self.manifestObject:
            self.downloadManifestByImageResourceUri(resourceUri)
        if resourceUri != self.manifestObject.identifier:
            raise InputException('Given image ID [%s] does not match to downloaded [%s]' %
                                    (resourceUri, self.manifestObject.identifier))

    def downloadManifestByImageResourceUri(self, imageId):
        tempMetadataFilename = tempfile.mktemp()
        try:
            self.manifestObject = self._getManifest(imageId, tempMetadataFilename)
        finally:
            try:
                os.unlink(tempMetadataFilename)
            except:
                pass

    def download(self, uri):
        '''uri is the full resource uri uniquely identifying
           a single manifest entry'''
        tempMetadataFilename = tempfile.mktemp()

        manifestInfo = self._getManifest(uri, tempMetadataFilename)

        tempImageFilename = ''
        for location in manifestInfo.locations:
            self._printDetail('Looking for image: %s' % location)
            try:
                tempImageFilename = self._downloadImage(location)
                break
            except KeyboardInterrupt:
                raise
            except:
                pass

        if not os.path.exists(tempImageFilename):
            raise InputException('Failed to find image matching image resource uri: %s' % uri)

        self._verifySignature(tempImageFilename, tempMetadataFilename)

        tempImageFilename = self._inflateImage(tempImageFilename)

        self._verifyHash(tempImageFilename, manifestInfo.sha1)

        shutil.copy2(tempImageFilename, self.localImageFilename)

        os.remove(tempImageFilename)
        os.remove(tempMetadataFilename)

        return self.localImageFilename

    def _loadDom(self, filename):
        file = open(filename).read()
        dom = etree.fromstring(file)
        return dom

    def _downloadImage(self, url):
        compressionExtension = self._extractCompressionExtension(url)

        localFilename = tempfile.mktemp()
        localImageName = localFilename + compressionExtension

        Util.wget(url, localImageName)

        return localImageName

    def _extractCompressionExtension(self, url):
        compression = url.split('.')[-1]

        if compression in Compressor.compressionFormats:
            return '.' + compression
        else:
            return ''

    def _download(self, url, localFilename):
        try:
            return Util.wget(url, localFilename)
        except urllib2.URLError, ex:
            raise InputException('Failed to download: %s, with detail: %s' % (url, str(ex)))

    def _verifySignature(self, imageFilename, metadataFilename):
        signator = Signator(metadataFilename, self.configHolder)
        res = signator.validate()
        if res:
            raise ExecutionException('Failed to validate metadata file')

    def _inflateImage(self, imageFilename):
        extension = self._extractCompressionExtension(imageFilename)
        inflatedFilename = imageFilename
        if extension:
            self._printDetail('Inflating image %s' % imageFilename)
            Compressor.inflate(imageFilename)
            inflatedFilename = imageFilename[:-len(extension)]

        return inflatedFilename

    def _verifyHash(self, imageFilename, hashFromManifest):

        data = open(imageFilename).read()
        sha1 = hashlib.sha1()
        sha1.update(data)

        imageHash = sha1.hexdigest()

        if imageHash != hashFromManifest:
            raise ValidationException('Manifest hash code different from downloaded image: image=%s, matadata=%s' % (imageHash, hashFromManifest))

    def _printDetail(self, message):
        Util.printDetail(message, self.verboseLevel, 1)

<|endoftext|>


<commit_before>from odoo.tests.common import SingleTransactionCase, get_db_name, at_install, post_install
from ..models.res_users import WEBSITE_REFS
from odoo.api import Environment

db_name = get_db_name()

@at_install(True)
@post_install(True)
class TestWebsiteMultiCompany(SingleTransactionCase):

    def _test_website_price_difference_is_accessible(self, env):
        website = env.ref(WEBSITE_REFS[0])
        products = env['product.template'].search(
            [('company_id', '=', website.company_id.id)] +
            website.sale_product_domain()
        )
        product = products[0]
        # make sure, it does not throw exception
        product.website_price_difference  # pylint: disable=pointless-statement

    def test_website_price_difference_is_accessible_for_demo_user(self):
        uid = self.registry['res.users'].authenticate(db_name, 'demo', 'demo', {})
        with self.cursor() as cr:
            env = Environment(cr, uid, {})
            self._test_website_price_difference_is_accessible(env)

    def test_website_price_difference_is_accessible_for_public_user(self):
        self._test_website_price_difference_is_accessible(self.env(user=self.browse_ref('base.public_user')))
<commit_msg>:green_heart: website_multi_company_demo: lint<commit_after>from odoo.tests.common import SingleTransactionCase, get_db_name, at_install, post_install
from ..models.res_users import WEBSITE_REFS
from odoo.api import Environment

db_name = get_db_name()

@at_install(True)
@post_install(True)
class TestWebsiteMultiCompany(SingleTransactionCase):

    def _test_website_price_difference_is_accessible(self, env):
        website = env.ref(WEBSITE_REFS[0])
        products = env['product.template'].search(
            [('company_id', '=', website.company_id.id)] + website.sale_product_domain()
        )
        product = products[0]
        # make sure, it does not throw exception
        product.website_price_difference  # pylint: disable=pointless-statement

    def test_website_price_difference_is_accessible_for_demo_user(self):
        uid = self.registry['res.users'].authenticate(db_name, 'demo', 'demo', {})
        with self.cursor() as cr:
            env = Environment(cr, uid, {})
            self._test_website_price_difference_is_accessible(env)

    def test_website_price_difference_is_accessible_for_public_user(self):
        self._test_website_price_difference_is_accessible(self.env(user=self.browse_ref('base.public_user')))
<|endoftext|>


<commit_before>from django.core.management import call_command
from django.test import TestCase

from corehq.motech.dhis2.repeaters import (
    SQLDhis2EntityRepeater,
    SQLDhis2Repeater,
)
from corehq.motech.fhir.repeaters import SQLFHIRRepeater
from corehq.motech.models import ConnectionSettings
from corehq.motech.repeaters.dbaccessors import delete_all_repeaters
from corehq.motech.repeaters.expression.repeaters import (
    SQLCaseExpressionRepeater,
)
from corehq.motech.repeaters.models import (
    Repeater,
    SQLAppStructureRepeater,
    SQLCaseRepeater,
    SQLCreateCaseRepeater,
    SQLDataRegistryCaseUpdateRepeater,
    SQLFormRepeater,
    SQLLocationRepeater,
    SQLReferCaseRepeater,
    SQLRepeater,
    SQLShortFormRepeater,
    SQLUpdateCaseRepeater,
    SQLUserRepeater,
)
from corehq.motech.openmrs.repeaters import SQLOpenmrsRepeater

from .data.repeaters import repeater_test_data

class TestMigrationCommand(TestCase):

    @classmethod
    def setUpClass(cls) -> None:
        cls.conn = ConnectionSettings(id=1, url="http://url.com", domain='rtest')
        cls.conn.save()
        cls.couch_repeaters = []
        for r in repeater_test_data:
            r = Repeater.wrap(r)
            r.save(sync_to_sql=False)
            cls.couch_repeaters.append(r)
        return super().setUpClass()

    @classmethod
    def tearDownClass(cls) -> None:
        delete_all_repeaters()
        return super().tearDownClass()

    def test_case_repeater_docs_are_migrated(self):
        call_command('migrate_caserepeater')
        self._assert_repeaters_equality(SQLCaseRepeater, "CaseRepeater")
        call_command('migrate_formrepeater')
        self._assert_repeaters_equality(SQLFormRepeater, "FormRepeater")
        call_command('migrate_shortformrepeater')
        self._assert_repeaters_equality(SQLShortFormRepeater, "ShortFormRepeater")
        call_command('migrate_createcaserepeater')
        self._assert_repeaters_equality(SQLCreateCaseRepeater, "CreateCaseRepeater")
        call_command('migrate_refercaserrepeater')
        self._assert_repeaters_equality(SQLReferCaseRepeater, "ReferCaseRepeater")
        call_command('migrate_dhis2repeater')
        self._assert_repeaters_equality(SQLDhis2Repeater, "Dhis2Repeater")
        call_command('migrate_userrepeater')
        self._assert_repeaters_equality(SQLUserRepeater, "UserRepeater")
        call_command('migrate_fhirrepeater')
        self._assert_repeaters_equality(SQLFHIRRepeater, "FHIRRepeater")
        call_command('migrate_appstructurerepeater')
        self._assert_repeaters_equality(SQLAppStructureRepeater, "AppStructureRepeater")
        call_command('migrate_caseexpressionrepeater')
        self._assert_repeaters_equality(SQLCaseExpressionRepeater, "CaseExpressionRepeater")
        call_command('migrate_dataregistrycaseupdaterepeater')
        self._assert_repeaters_equality(SQLDataRegistryCaseUpdateRepeater, "DataRegistryCaseUpdateRepeater")
        call_command('migrate_dhis2entityrepeater')
        self._assert_repeaters_equality(SQLDhis2EntityRepeater, "Dhis2EntityRepeater")
        call_command('migrate_openmrsrepeater')
        self._assert_repeaters_equality(SQLOpenmrsRepeater, "OpenmrsRepeater")
        call_command('migrate_locationrepeater')
        self._assert_repeaters_equality(SQLLocationRepeater, "LocationRepeater")
        call_command('migrate_updatecaserepeater')
        self._assert_repeaters_equality(SQLUpdateCaseRepeater, "UpdateCaseRepeater")

        # test for count
        self.assertEqual(SQLRepeater.objects.count(), len(self.couch_repeaters))

    def _assert_repeaters_equality(self, sql_class, doc_type):
        sql_ids = set(sql_class.objects.all().values_list('repeater_id', flat=True))
        couch_ids = {r._id for r in self._get_repeater_objects(doc_type)}
        self.assertEqual(len(couch_ids), 2)
        self.assertEqual(len(sql_ids), 2)
        self.assertCountEqual(sql_ids, couch_ids)
        self.assertEqual(sql_ids, couch_ids)

    def _get_repeater_objects(self, repeater_type):
        return [r for r in self.couch_repeaters if r.doc_type == repeater_type]
<commit_msg>add test for equality of config attrs<commit_after>from django.core.management import call_command
from django.test import TestCase

from corehq.motech.dhis2.repeaters import (
    SQLDhis2EntityRepeater,
    SQLDhis2Repeater,
)
from corehq.motech.fhir.repeaters import SQLFHIRRepeater
from corehq.motech.models import ConnectionSettings
from corehq.motech.repeaters.dbaccessors import delete_all_repeaters
from corehq.motech.repeaters.expression.repeaters import (
    SQLCaseExpressionRepeater,
)
from corehq.motech.repeaters.models import (
    Repeater,
    SQLAppStructureRepeater,
    SQLCaseRepeater,
    SQLCreateCaseRepeater,
    SQLDataRegistryCaseUpdateRepeater,
    SQLFormRepeater,
    SQLLocationRepeater,
    SQLReferCaseRepeater,
    SQLRepeater,
    SQLShortFormRepeater,
    SQLUpdateCaseRepeater,
    SQLUserRepeater,
)
from corehq.motech.openmrs.repeaters import SQLOpenmrsRepeater

from .data.repeaters import repeater_test_data

class TestMigrationCommand(TestCase):

    @classmethod
    def setUpClass(cls) -> None:
        cls.conn = ConnectionSettings(id=1, url="http://url.com", domain='rtest')
        cls.conn.save()
        cls.couch_repeaters = []
        for r in repeater_test_data:
            r = Repeater.wrap(r)
            r.save(sync_to_sql=False)
            cls.couch_repeaters.append(r)
        return super().setUpClass()

    @classmethod
    def tearDownClass(cls) -> None:
        delete_all_repeaters()
        return super().tearDownClass()

    def test_case_repeater_docs_are_migrated(self):
        call_command('migrate_caserepeater')
        self._assert_repeaters_equality(SQLCaseRepeater, "CaseRepeater")
        call_command('migrate_formrepeater')
        self._assert_repeaters_equality(SQLFormRepeater, "FormRepeater")
        call_command('migrate_shortformrepeater')
        self._assert_repeaters_equality(SQLShortFormRepeater, "ShortFormRepeater")
        call_command('migrate_createcaserepeater')
        self._assert_repeaters_equality(SQLCreateCaseRepeater, "CreateCaseRepeater")
        call_command('migrate_refercaserrepeater')
        self._assert_repeaters_equality(SQLReferCaseRepeater, "ReferCaseRepeater")
        call_command('migrate_dhis2repeater')
        self._assert_repeaters_equality(SQLDhis2Repeater, "Dhis2Repeater")
        call_command('migrate_userrepeater')
        self._assert_repeaters_equality(SQLUserRepeater, "UserRepeater")
        call_command('migrate_fhirrepeater')
        self._assert_repeaters_equality(SQLFHIRRepeater, "FHIRRepeater")
        call_command('migrate_appstructurerepeater')
        self._assert_repeaters_equality(SQLAppStructureRepeater, "AppStructureRepeater")
        call_command('migrate_caseexpressionrepeater')
        self._assert_repeaters_equality(SQLCaseExpressionRepeater, "CaseExpressionRepeater")
        call_command('migrate_dataregistrycaseupdaterepeater')
        self._assert_repeaters_equality(SQLDataRegistryCaseUpdateRepeater, "DataRegistryCaseUpdateRepeater")
        call_command('migrate_dhis2entityrepeater')
        self._assert_repeaters_equality(SQLDhis2EntityRepeater, "Dhis2EntityRepeater")
        call_command('migrate_openmrsrepeater')
        self._assert_repeaters_equality(SQLOpenmrsRepeater, "OpenmrsRepeater")
        call_command('migrate_locationrepeater')
        self._assert_repeaters_equality(SQLLocationRepeater, "LocationRepeater")
        call_command('migrate_updatecaserepeater')
        self._assert_repeaters_equality(SQLUpdateCaseRepeater, "UpdateCaseRepeater")

        # test for count
        self.assertEqual(SQLRepeater.objects.count(), len(self.couch_repeaters))

    def _assert_repeaters_equality(self, sql_class, doc_type):
        sql_ids = set(sql_class.objects.all().values_list('repeater_id', flat=True))
        couch_ids = {r._id for r in self._get_repeater_objects(doc_type)}
        self.assertEqual(len(couch_ids), 2)
        self.assertEqual(len(sql_ids), 2)
        self.assertCountEqual(sql_ids, couch_ids)
        self.assertEqual(sql_ids, couch_ids)

    def test_equality_of_config_attrs(self):
        call_command('migrate_dhis2repeater')
        call_command('migrate_openmrsrepeater')
        call_command('migrate_caseexpressionrepeater')
        dhsi2_objects = self._get_repeater_objects('Dhis2Repeater')
        openmrs_objects = self._get_repeater_objects('OpenmrsRepeater')
        caseexpression_objects = self._get_repeater_objects('CaseExpressionRepeater')

        for obj in dhsi2_objects:
            sql_obj = SQLDhis2Repeater.objects.get(repeater_id=obj._id)
            self.assertEqual(sql_obj.dhis2_config, obj.dhis2_config.to_json())

        for obj in openmrs_objects:
            sql_obj = SQLOpenmrsRepeater.objects.get(repeater_id=obj._id)
            self.assertEqual(sql_obj.openmrs_config, obj.openmrs_config.to_json())
            self.assertEqual(sql_obj.atom_feed_status, obj.atom_feed_status)

        for obj in caseexpression_objects:
            sql_obj = SQLCaseExpressionRepeater.objects.get(repeater_id=obj._id)
            self.assertEqual(sql_obj.configured_filter, obj.configured_filter)
            self.assertEqual(sql_obj.configured_expression, obj.configured_expression)

    def _get_repeater_objects(self, repeater_type):
        return [r for r in self.couch_repeaters if r.doc_type == repeater_type]
<|endoftext|>


<commit_before>from __future__ import print_function
import random
import unittest
import operator

from cpy2py.utility.compat import rangex
from cpy2py.utility.thread_tools import ThreadGuard

class TestThreadGuard(unittest.TestCase):
    def test_counter(self):
        counter = 1
        tg_counter = ThreadGuard(str(counter))
        for op in [
            '__add__', '__sub__', '__mul__', '__div__', '__truediv__', '__floordiv__', '__pow__', '__mod__',
            '__eq__', '__ne__', '__lt__', '__gt__', '__le__', '__ge__',
            '__and__', '__or__', '__xor__',
            '__lshift__', '__rshift__',
        ]:
            try:
                ops = getattr(operator, op)
            except AttributeError:
                continue
            else:
                for val in (random.randint(-100, 100) for _ in rangex(5)):
                    if val == 0:  # zero division
                        val = 1
                    if op in ('__lshift__', '__rshift__'):
                        val = abs(val)
                    res = ops(counter, val)
                    tg_res = ops(tg_counter, val)
                    self.assertEqual(res, tg_res, '%s %s %s %s' % (res, tg_res, op, val))
                    rres = ops(val, counter)
                    rtg_res = ops(val, tg_counter)
                    self.assertEqual(rres, rtg_res, '%s %s %s %s (reflected)' % (rres, rtg_res, op, val))
<commit_msg>added inplace and conversion tests<commit_after>from __future__ import print_function
import random
import unittest
import operator

from cpy2py.utility.compat import rangex
from cpy2py.utility.thread_tools import ThreadGuard

class TestThreadGuard(unittest.TestCase):
    def test_counter_ops(self):
        counter = 100
        tg_counter = ThreadGuard(str(counter))
        for op in [
            '__add__', '__sub__', '__mul__', '__div__', '__truediv__', '__floordiv__', '__pow__', '__mod__',
            '__eq__', '__ne__', '__lt__', '__gt__', '__le__', '__ge__',
            '__and__', '__or__', '__xor__',
            '__lshift__', '__rshift__',
        ]:
            try:
                ops = getattr(operator, op)
            except AttributeError:
                continue
            else:
                for val in (random.randint(-100, 100) for _ in rangex(5)):
                    if val == 0:  # zero division
                        val = 1
                    if op in ('__lshift__', '__rshift__'):
                        val = abs(val)
                    res = ops(counter, val)
                    tg_res = ops(tg_counter, val)
                    self.assertEqual(res, tg_res, '%s %s %s %s' % (res, tg_res, op, val))
                    rres = ops(val, counter)
                    rtg_res = ops(val, tg_counter)
                    self.assertEqual(rres, rtg_res, '%s %s %s %s (reflected)' % (rres, rtg_res, op, val))

    def test_counter_inplace_ops(self):
        for op in [
            '__iadd__', '__isub__', '__imul__', '__idiv__', '__itruediv__', '__ifloordiv__', '__ipow__', '__imod__',
            '__iand__', '__ior__', '__ixor__',
            '__ilshift__', '__irshift__',
        ]:
            try:
                ops = getattr(operator, op)
            except AttributeError:
                continue
            else:
                counter = 100
                tg_counter = ThreadGuard(str(counter))
                for val in (random.randint(-100, 100) for _ in rangex(5)):
                    if val == 0:  # zero division
                        val = 1
                    if op in ('__ilshift__', '__irshift__', '__ipow__'):
                        val = random.randint(0, 5)
                    counter = ops(counter, val)
                    tg_counter = ops(tg_counter, val)
                    self.assertEqual(counter, tg_counter, '%s %s %s %s' % (counter, tg_counter, op, val))

    def test_counter_convert(self):
        counter = 100
        tg_counter = ThreadGuard(str(counter))
        for conversion in (float, int, complex, bool, round, str):
            self.assertEqual(conversion(counter), conversion(tg_counter))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Implementation of http://elementalselenium.com/tips/8-download-a-file-revisited
"""

import unittest
from selenium import webdriver
import httplib # Use http.client if using Python 3.x.x

class FileDownloadRevisited(unittest.TestCase):

    def setUp(self):
        self.driver = webdriver.Firefox()

    def tearDown(self):
        self.driver.quit()

    def test_example_1(self):
        driver = self.driver
        driver.get('http://the-internet.herokuapp.com/download')
        download_link = driver.find_element_by_css_selector('.example a').get_attribute('href')

        connection = httplib.HTTPConnection('the-internet.herokuapp.com')
        connection.request('HEAD', download_link)
        response = connection.getresponse()
        content_type = response.getheader('Content-type')
        content_length = response.getheader('Content-length')

        assert content_type == 'application/octet-stream'
        assert content_length > 0

if __name__ == "__main__":
    unittest.main()
<commit_msg>Updated for Python 3<commit_after># -*- coding: utf-8 -*-
"""
Implementation of http://elementalselenium.com/tips/8-download-a-file-revisited
"""

import unittest
from selenium import webdriver
import http.client # Use httplib if using Python 2.x.x

class FileDownloadRevisited(unittest.TestCase):

    def setUp(self):
        self.driver = webdriver.Firefox()

    def tearDown(self):
        self.driver.quit()

    def test_example_1(self):
        driver = self.driver
        driver.get('http://the-internet.herokuapp.com/download')
        download_link = driver.find_element_by_css_selector('.example a').get_attribute('href')

        connection = http.client.HTTPConnection('the-internet.herokuapp.com')
        connection.request('HEAD', download_link)
        response = connection.getresponse()
        connection.close()
        content_type = response.getheader('Content-type')
        content_length = int(response.getheader('Content-length'))

        assert content_type == 'application/octet-stream'
        assert content_length > 0

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