<commit_before>import io
import os
import threading
import unittest

from jobmon.protocol import *

class TestProtocol:
    def test_events(self):
        events = (EVENT_STARTJOB, EVENT_STOPJOB, EVENT_RESTARTJOB,
                EVENT_TERMINATE)

        proto_read, proto_write = self.make_protocol()

        try:
            for event_code in events:
                event = Event('some_job', event_code)
                proto_write.send(event)

                out_event = proto_read.recv()
                self.assertEqual(out_event, event)
        finally:
            self.cleanup_protocol(proto_read, proto_write)

    def test_commands(self):
        commands = (CMD_START, CMD_STOP, CMD_STATUS, CMD_JOB_LIST, CMD_QUIT)
        proto_read, proto_write = self.make_protocol()

        try:
            for command_code in commands:
                command = Command('some_job', command_code)
                proto_write.send(command)

                out_command = proto_read.recv()
                self.assertEqual(out_command, command)
        finally:
            self.cleanup_protocol(proto_read, proto_write)

    def test_resonses(self):
        responses = (SuccessResponse('some_job'),
                FailureResponse('some_job', ERR_NO_SUCH_JOB),
                FailureResponse('some_job', ERR_JOB_STARTED),
                FailureResponse('some_job', ERR_JOB_STOPPED),
                StatusResponse('some_job', True),
                StatusResponse('some_job', False),
                JobListResponse({'a': True, 'b': False}))

        proto_read, proto_write = self.make_protocol()
        try:
            for response in responses:
                proto_write.send(response)

                out_response = proto_read.recv()
                self.assertEqual(out_response, response)
        finally:
            self.cleanup_protocol(proto_read, proto_write)

class TestProtocolFile(TestProtocol, unittest.TestCase):
    def make_protocol(self):
        reader, writer = os.pipe()

        protocol_reader = ProtocolFile(os.fdopen(reader, 'rb'))
        protocol_writer = ProtocolFile(os.fdopen(writer, 'wb'))
        return protocol_reader, protocol_writer

    def cleanup_protocol(self, reader, writer):
        reader.close()
        writer.close()

PORT = 9999

class TestProtocolStreamSocket(TestProtocol, unittest.TestCase):
    class ServerThread(threading.Thread):
        def __init__(self):
            super().__init__()
            self.wait_for_start = threading.Event()
            self.wait_for_client = threading.Event()
            self.initialize_exit = threading.Event()
            self.wait_for_exit = threading.Event()

            self.client = None

        def run(self):
            server = socket.socket()
            server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

            server.bind(('localhost', PORT))
            server.listen(1)
            self.wait_for_start.set()

            self.client, _ = server.accept()
            self.wait_for_client.set()

            self.initialize_exit.wait()

            self.client.close()
            server.close()

            self.wait_for_exit.set()

        def wait_start(self):
            self.wait_for_start.wait()

        def wait_client(self):
            self.wait_for_client.wait()

        def exit(self):
            self.initialize_exit.set()
            self.wait_for_exit.wait()

    def make_protocol(self):
        self._thread = self.ServerThread()
        self._thread.start()
        self._thread.wait_start()

        writer = socket.socket()
        writer.connect(('localhost', PORT))
        self._thread.wait_client()

        reader = self._thread.client

        return ProtocolStreamSocket(reader), ProtocolStreamSocket(writer)

    def cleanup_protocol(self, reader, writer):
        self._thread.exit()
        writer.close()
        # The reader is already closed, since the server thread is done

class TestProtocolDatagramSocket(TestProtocol, unittest.TestCase):
    class ServerThread(threading.Thread):
        def __init__(self):
            super().__init__()
            self.wait_for_start = threading.Event()
            self.initialize_exit = threading.Event()
            self.wait_for_exit = threading.Event()

            self.server = None

        def run(self):
            self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.server.bind(('localhost', PORT))

            self.wait_for_start.set()

            self.initialize_exit.wait()

            self.server.close()

            self.wait_for_exit.set()

        def wait_start(self):
            self.wait_for_start.wait()

        def exit(self):
            self.initialize_exit.set()
            self.wait_for_exit.wait()

    def make_protocol(self):
        self._thread = self.ServerThread()
        self._thread.start()
        self._thread.wait_start()

        _writer = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        writer = ProtocolDatagramSocket(_writer, ('localhost', PORT))

        _reader = self._thread.server
        reader = ProtocolDatagramSocket(_reader, None)

        return reader, writer

    def cleanup_protocol(self, reader, writer):
        self._thread.exit()
        writer.close()
        # The reader is already closed, since the server thread is done
<commit_msg>Fix an importing issue with the tests<commit_after>import io
import os
import socket
import threading
import unittest

from jobmon.protocol import *

class TestProtocol:
    def test_events(self):
        events = (EVENT_STARTJOB, EVENT_STOPJOB, EVENT_RESTARTJOB,
                EVENT_TERMINATE)

        proto_read, proto_write = self.make_protocol()

        try:
            for event_code in events:
                event = Event('some_job', event_code)
                proto_write.send(event)

                out_event = proto_read.recv()
                self.assertEqual(out_event, event)
        finally:
            self.cleanup_protocol(proto_read, proto_write)

    def test_commands(self):
        commands = (CMD_START, CMD_STOP, CMD_STATUS, CMD_JOB_LIST, CMD_QUIT)
        proto_read, proto_write = self.make_protocol()

        try:
            for command_code in commands:
                command = Command('some_job', command_code)
                proto_write.send(command)

                out_command = proto_read.recv()
                self.assertEqual(out_command, command)
        finally:
            self.cleanup_protocol(proto_read, proto_write)

    def test_resonses(self):
        responses = (SuccessResponse('some_job'),
                FailureResponse('some_job', ERR_NO_SUCH_JOB),
                FailureResponse('some_job', ERR_JOB_STARTED),
                FailureResponse('some_job', ERR_JOB_STOPPED),
                StatusResponse('some_job', True),
                StatusResponse('some_job', False),
                JobListResponse({'a': True, 'b': False}))

        proto_read, proto_write = self.make_protocol()
        try:
            for response in responses:
                proto_write.send(response)

                out_response = proto_read.recv()
                self.assertEqual(out_response, response)
        finally:
            self.cleanup_protocol(proto_read, proto_write)

class TestProtocolFile(TestProtocol, unittest.TestCase):
    def make_protocol(self):
        reader, writer = os.pipe()

        protocol_reader = ProtocolFile(os.fdopen(reader, 'rb'))
        protocol_writer = ProtocolFile(os.fdopen(writer, 'wb'))
        return protocol_reader, protocol_writer

    def cleanup_protocol(self, reader, writer):
        reader.close()
        writer.close()

PORT = 9999

class TestProtocolStreamSocket(TestProtocol, unittest.TestCase):
    class ServerThread(threading.Thread):
        def __init__(self):
            super().__init__()
            self.wait_for_start = threading.Event()
            self.wait_for_client = threading.Event()
            self.initialize_exit = threading.Event()
            self.wait_for_exit = threading.Event()

            self.client = None

        def run(self):
            server = socket.socket()
            server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

            server.bind(('localhost', PORT))
            server.listen(1)
            self.wait_for_start.set()

            self.client, _ = server.accept()
            self.wait_for_client.set()

            self.initialize_exit.wait()

            self.client.close()
            server.close()

            self.wait_for_exit.set()

        def wait_start(self):
            self.wait_for_start.wait()

        def wait_client(self):
            self.wait_for_client.wait()

        def exit(self):
            self.initialize_exit.set()
            self.wait_for_exit.wait()

    def make_protocol(self):
        self._thread = self.ServerThread()
        self._thread.start()
        self._thread.wait_start()

        writer = socket.socket()
        writer.connect(('localhost', PORT))
        self._thread.wait_client()

        reader = self._thread.client

        return ProtocolStreamSocket(reader), ProtocolStreamSocket(writer)

    def cleanup_protocol(self, reader, writer):
        self._thread.exit()
        writer.close()
        # The reader is already closed, since the server thread is done

class TestProtocolDatagramSocket(TestProtocol, unittest.TestCase):
    class ServerThread(threading.Thread):
        def __init__(self):
            super().__init__()
            self.wait_for_start = threading.Event()
            self.initialize_exit = threading.Event()
            self.wait_for_exit = threading.Event()

            self.server = None

        def run(self):
            self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.server.bind(('localhost', PORT))

            self.wait_for_start.set()

            self.initialize_exit.wait()

            self.server.close()

            self.wait_for_exit.set()

        def wait_start(self):
            self.wait_for_start.wait()

        def exit(self):
            self.initialize_exit.set()
            self.wait_for_exit.wait()

    def make_protocol(self):
        self._thread = self.ServerThread()
        self._thread.start()
        self._thread.wait_start()

        _writer = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        writer = ProtocolDatagramSocket(_writer, ('localhost', PORT))

        _reader = self._thread.server
        reader = ProtocolDatagramSocket(_reader, None)

        return reader, writer

    def cleanup_protocol(self, reader, writer):
        self._thread.exit()
        writer.close()
        # The reader is already closed, since the server thread is done
<|endoftext|>


<commit_before># ===============================================================================
# Copyright 2016 ross
#
# 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.
# ===============================================================================

from __future__ import absolute_import
from traits.api import Str, List, HasTraits, Dict, Any
from traitsui.api import View, Item, EnumEditor

# ============= standard library imports ========================
# ============= local library imports  ==========================

class AddAnalysisGroupView(HasTraits):
    name = Str
    project = Any
    projects = Dict

    def save(self, ans, db):
        append = False

        gdb = db.get_analysis_groups_by_name(self.name, self.project)
        ok = True
        if gdb:
            gdb = gdb[-1]
            if db.confirmation_dialog('"{}" already exists? Would you like to append your selection'.format(gdb.name)):
                append = True
            else:
                ok = False

        if append:
            db.append_analysis_group(gdb, ans)
        elif ok:
            db.add_analysis_group(ans, self.name, self.project)

    def traits_view(self):
        v = View(Item('name'),
                 Item('project', editor=EnumEditor(name='projects')),
                 resizable=True,
                 buttons=['OK', 'Cancel'],
                 title='Add Analysis Group')
        return v

# ============= EOF =============================================
<commit_msg>addressed issue #1450<commit_after># ===============================================================================
# Copyright 2016 ross
#
# 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.
# ===============================================================================

from __future__ import absolute_import
from traits.api import Str, List, HasTraits, Dict, Any
from traitsui.api import View, Item, EnumEditor

# ============= standard library imports ========================
# ============= local library imports  ==========================

class AddAnalysisGroupView(HasTraits):
    name = Str
    project = Any
    projects = Dict

    def save(self, ans, db):
        append = False

        gdb = db.get_analysis_groups_by_name(self.name, self.project)
        ok = True
        if gdb:
            gdb = gdb[-1]
            if db.confirmation_dialog('"{}" already exists? Would you like to append your selection'.format(gdb.name)):
                append = True
            else:
                ok = False

        if append:
            db.append_analysis_group(gdb, ans)
        elif ok:
            db.add_analysis_group(ans, self.name, self.project.name)

    def traits_view(self):
        v = View(Item('name'),
                 Item('project', editor=EnumEditor(name='projects')),
                 resizable=True,
                 buttons=['OK', 'Cancel'],
                 title='Add Analysis Group')
        return v

# ============= EOF =============================================
<|endoftext|>


<commit_before>import cis_profile
import cis_publisher
import boto3
import botocore
import os
import logging
import json
import time
import requests
from auth0.v3.authentication import GetToken
from auth0.v3.management import Auth0
from datetime import datetime, timezone, timedelta
from traceback import format_exc

logger = logging.getLogger(__name__)

class Auth0Publisher:
    def __init__(self, context={}):
        self.secret_manager = cis_publisher.secret.Manager()
        self.context = context
        self.report = None
        self.config = cis_publisher.common.get_config()
        # Only fields we care about for the user entries
        # auth0 field->cis field map
        self.az_cis_fields = {
            "created_at": "created",
            "updated_at": "last_modified",
            "given_name": "first_name",
            "family_name": "last_name",
            "name": None,
            "nickname": None,
            "picture": "picture",
            "user_id": "user_id",
            "email": "primary_email",
        }

    def publish(self, user_ids=None, chunk_size=30):
        """
        Glue to create or fetch cis_profile.User profiles for this publisher
        Then pass everything over to the Publisher class
        None, ALL profiles are sent.
        @user_ids: list of str - user ids to publish. If None, all users are published.
        @chunk_size: int when no user_id is selected, this is the size of the chunk/slice we'll create to divide the
        work between function calls (to self)
        """
        if user_ids is None:
            le = "All"
        else:
            le = len(user_ids)
        logger.info("Starting Auth0 Publisher [{} users]".format(le))

        # Should we fan-out processing to multiple function calls?
        if user_ids is None:
            self.fan_out(publisher, chunk_size)
        else:
            self.process(publisher, user_ids)

    def fetch_az_users(self):
        """
        Fetches ALL valid users from auth0'z database
        Returns list of user attributes
        """
        az_api_url = self.config("auth0_api_url", namespace="cis", default="auth-dev.mozilla.auth0.com")
        az_client_id = (self.secret_manager.secret("client_id"),)
        az_client_secret = (self.secret_manager.secret("client_secret"),)
        az_fields = self.az_cis_fields.keys()

        # Build the connection query (excludes LDAP)
        az_blacklist_connections = ["Mozilla-LDAP", "Mozilla-LDAP-Dev"]
        # Ignore non-verified `email` (such as unfinished passwordless flows) as we don't consider these to be valid
        # users
        az_query = "email_verified:true AND NOT ("
        t = ""
        for azc in az_connections:
            az_query = azquery + t + 'identities.connection:"{}"'.format(az_connections)
            t = " OR "
        az_query = az_query + ")"

        logger.debug("About to get Auth0 user list")
        az_getter = GetToken(az_api_url)
        az_token = az_getter.client_credentials(as_client_id, az_client_secret, "https://{}/api/v2/".format(az_api_url))
        auth0 = Auth0(az_api_url, az_token["access_token"])

        # Query the entire thing
        p = 0
        user_list = []
        # This is an artificial upper limit of 100*99999 (per_page*page) i.e. 9 999 900 users max - just in case things
        # go wrong
        for p in range(0, 99999):
            tmp = auth0.users.list(page=p, per_page=100, fields=az_fields, q=az_query)
            if tmp == []:
                # stop when our page is empty
                logger.debug("Crawled {} pages from auth0 users API".format(p))
                break
            else:
                user_list.extend(tmp)
        logger.info("Received {} users from auth0".format(len(user_list)))
        return user_list

    def convert_az_users(self, az_users):
        """
        Convert a list of auth0 user fields to cis_profile Users
        @az_users list of dicts with user attributes
        Returns [cis_profile.Users]
        """
        profiles = []
        for u in az_users:
            p = cis_profile.User()
            for k, v in self.az_cis_fields.items():
                if k not in u.keys():
                    continue
                p.__dict__[v]["value"] = u[k]

            # Hack:
            # In case no names were passed, try to fix things up with what we have
            if p.__dict__["first_name"]["value"] is None:
                logger.debug("User {} has no name, asserting a name from other values".format(u["user_id"]))
                if "name" in u.keys():
                    p.__dict__["first_name"]["value"] = u["name"]
                elif "nickname" in u.keys():
                    p.__dict__["first_name"]["value"] = u["nickname"]
                else:
                    p.__dict__["first_name"]["value"] = u["email"].split("@")[0]
            profiles.append(p)
        return profiles

    def process(self, publisher, user_ids):
        """
        Process profiles and post them
        @publisher object the publisher object to operate on
        @user_ids list of user ids to process in this batch
        """

        profiles = self.convert_az_users(self.fetch_az_users())
        logger.info("Processing {} profiles".format(len(profiles)))
        publisher.profiles = profiles

        failures = []
        try:
            failures = publisher.post_all(user_ids=user_ids)
        except Exception as e:
            logger.error("Failed to post_all() profiles. Trace: {}".format(format_exc()))
            raise e

        if len(failures) > 0:
            logger.error("Failed to post {} profiles: {}".format(len(failures), failures))

    def fan_out(self, publisher, chunk_size):
        """
        Splices all users to process into chunks
        and self-invoke as many times as needed to complete all work in parallel lambda functions
        When self-invoking, this will effectively call self.process() instead of self.fan_out()
"]
        Note: chunk_size should never result in the invoke() argument to exceed 128KB (len(Payload.encode('utf-8') <
        128KB) as this is the maximum AWS Lambda payload size.

        @publisher object the cis_publisher object to operate on
        @chunk_size int size of the chunk to process
        """
        all_user_ids = []
        sliced = [all_user_ids[i : i + chunk_size] for i in range(0, len(all_user_ids), chunk_size)]
        logger.info(
            "No user_id selected. Creating slices of work, chunck size: {}, slices: {}, total users: {} and "
            "faning-out work to self".format(chunk_size, len(sliced), len(all_user_ids))
        )
        lambda_client = boto3.client("lambda")
        for s in sliced:
            lambda_client.invoke(FunctionName=self.context.function_name, InvocationType="Event", Payload=json.dumps(s))
            time.sleep(3)  # give api calls a chance, otherwise this storms resources

        logger.info("Exiting slicing function successfully")
<commit_msg>support passing user_ids to fetch from auth0 so that when specified you don't wait for all users to be passed back<commit_after>import cis_profile
import cis_publisher
import boto3
import botocore
import os
import logging
import json
import time
import requests
from auth0.v3.authentication import GetToken
from auth0.v3.management import Auth0
from datetime import datetime, timezone, timedelta
from traceback import format_exc

logger = logging.getLogger(__name__)

class Auth0Publisher:
    def __init__(self, context={}):
        self.secret_manager = cis_publisher.secret.Manager()
        self.context = context
        self.report = None
        self.config = cis_publisher.common.get_config()
        # Only fields we care about for the user entries
        # auth0 field->cis field map
        self.az_cis_fields = {
            "created_at": "created",
            "updated_at": "last_modified",
            "given_name": "first_name",
            "family_name": "last_name",
            "name": None,
            "nickname": None,
            "picture": "picture",
            "user_id": "user_id",
            "email": "primary_email",
        }

    def publish(self, user_ids=None, chunk_size=30):
        """
        Glue to create or fetch cis_profile.User profiles for this publisher
        Then pass everything over to the Publisher class
        None, ALL profiles are sent.
        @user_ids: list of str - user ids to publish. If None, all users are published.
        @chunk_size: int when no user_id is selected, this is the size of the chunk/slice we'll create to divide the
        work between function calls (to self)
        """
        if user_ids is None:
            le = "All"
        else:
            le = len(user_ids)
        logger.info("Starting Auth0 Publisher [{} users]".format(le))

        # Should we fan-out processing to multiple function calls?
        if user_ids is None:
            self.fan_out(publisher, chunk_size)
        else:
            self.process(publisher, user_ids)

    def fetch_az_users(self, user_ids=[]):
        """
        Fetches ALL valid users from auth0'z database
        Returns list of user attributes
        """
        az_api_url = self.config("auth0_api_url", namespace="cis", default="auth-dev.mozilla.auth0.com")
        az_client_id = (self.secret_manager.secret("client_id"),)
        az_client_secret = (self.secret_manager.secret("client_secret"),)
        az_fields = self.az_cis_fields.keys()

        # Build the connection query (excludes LDAP)
        az_blacklist_connections = ["Mozilla-LDAP", "Mozilla-LDAP-Dev"]
        # Ignore non-verified `email` (such as unfinished passwordless flows) as we don't consider these to be valid
        # users
        az_query = "email_verified:true AND NOT ("
        t = ""
        for azc in az_connections:
            az_query = azquery + t + 'identities.connection:"{}"'.format(az_connections)
            t = " OR "
        az_query = az_query + ")"

        # Build query for user_ids if some are specified (else it gets all of them)
        if user_ids:
            logger.info("Restriction auth0 user query to user_ids: {}".format(user_ids))
            t = ""
            az_query = az_query + "AND ("
            for u in user_ids:
                az_query = az_query + t + 'user_id:"{}"'.format(u)
                t = " OR "
            az_query = az_query + ")"

        logger.debug("About to get Auth0 user list")
        az_getter = GetToken(az_api_url)
        az_token = az_getter.client_credentials(as_client_id, az_client_secret, "https://{}/api/v2/".format(az_api_url))
        auth0 = Auth0(az_api_url, az_token["access_token"])

        # Query the entire thing
        p = 0
        user_list = []
        # This is an artificial upper limit of 100*99999 (per_page*page) i.e. 9 999 900 users max - just in case things
        # go wrong
        for p in range(0, 99999):
            tmp = auth0.users.list(page=p, per_page=100, fields=az_fields, q=az_query)
            if tmp == []:
                # stop when our page is empty
                logger.debug("Crawled {} pages from auth0 users API".format(p))
                break
            else:
                user_list.extend(tmp)
        logger.info("Received {} users from auth0".format(len(user_list)))
        return user_list

    def convert_az_users(self, az_users):
        """
        Convert a list of auth0 user fields to cis_profile Users
        @az_users list of dicts with user attributes
        Returns [cis_profile.Users]
        """
        profiles = []
        for u in az_users:
            p = cis_profile.User()
            for k, v in self.az_cis_fields.items():
                if k not in u.keys():
                    continue
                p.__dict__[v]["value"] = u[k]

            # Hack:
            # In case no names were passed, try to fix things up with what we have
            if p.__dict__["first_name"]["value"] is None:
                logger.debug("User {} has no name, asserting a name from other values".format(u["user_id"]))
                if "name" in u.keys():
                    p.__dict__["first_name"]["value"] = u["name"]
                elif "nickname" in u.keys():
                    p.__dict__["first_name"]["value"] = u["nickname"]
                else:
                    p.__dict__["first_name"]["value"] = u["email"].split("@")[0]
            profiles.append(p)
        return profiles

    def process(self, publisher, user_ids):
        """
        Process profiles and post them
        @publisher object the publisher object to operate on
        @user_ids list of user ids to process in this batch
        """

        profiles = self.convert_az_users(self.fetch_az_users(user_ids))
        logger.info("Processing {} profiles".format(len(profiles)))
        publisher.profiles = profiles

        failures = []
        try:
            failures = publisher.post_all(user_ids=user_ids)
        except Exception as e:
            logger.error("Failed to post_all() profiles. Trace: {}".format(format_exc()))
            raise e

        if len(failures) > 0:
            logger.error("Failed to post {} profiles: {}".format(len(failures), failures))

    def fan_out(self, publisher, chunk_size):
        """
        Splices all users to process into chunks
        and self-invoke as many times as needed to complete all work in parallel lambda functions
        When self-invoking, this will effectively call self.process() instead of self.fan_out()
"]
        Note: chunk_size should never result in the invoke() argument to exceed 128KB (len(Payload.encode('utf-8') <
        128KB) as this is the maximum AWS Lambda payload size.

        @publisher object the cis_publisher object to operate on
        @chunk_size int size of the chunk to process
        """
        all_user_ids = []
        sliced = [all_user_ids[i : i + chunk_size] for i in range(0, len(all_user_ids), chunk_size)]
        logger.info(
            "No user_id selected. Creating slices of work, chunck size: {}, slices: {}, total users: {} and "
            "faning-out work to self".format(chunk_size, len(sliced), len(all_user_ids))
        )
        lambda_client = boto3.client("lambda")
        for s in sliced:
            lambda_client.invoke(FunctionName=self.context.function_name, InvocationType="Event", Payload=json.dumps(s))
            time.sleep(3)  # give api calls a chance, otherwise this storms resources

        logger.info("Exiting slicing function successfully")
<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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
#

from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function

import socket
import binascii

#
#
class PolicyError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return str(self.value)

def is_ipv6_enabled():
    """
    Returns true if IPV6 is enabled, false otherwise
    """
    ipv6_enabled = True
    try:
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.bind(('::1', 0))
    except Exception as e:
        ipv6_enabled = False

    return ipv6_enabled

class HostStruct(object):
    """
    HostStruct represents a single, binary socket address from getaddrinfo
        - name     : name given to constructor; numeric IP or host name
        - saddr    : net name resolved by getaddrinfo; numeric IP
        - family   : saddr.family; int
        - binary   : saddr packed binary address; binary string
    """
    families = [socket.AF_INET]
    famnames = ["IPv4"]
    if is_ipv6_enabled():
        families.append(socket.AF_INET6)
        famnames.append("IPv6")

    def __init__(self, hostname):
        """
        Given a host name text string, return the socket info for it.
        @param[in] hostname host IP address to parse
        """
        try:
            res = socket.getaddrinfo(hostname, 0, socket.AF_UNSPEC, socket.SOCK_STREAM)
            if len(res) == 0:
                raise PolicyError("HostStruct: '%s' did not resolve to an IP address" % hostname)
            foundFirst = False
            saddr = ""
            sfamily = socket.AF_UNSPEC
            for i0 in range(0, len(res)):
                family, dum0, dum1, dum2, sockaddr = res[i0]
                if not foundFirst:
                    if family in self.families:
                        saddr = sockaddr[0]
                        sfamily = family
                        foundFirst = True
                else:
                    if family in self.families:
                        if not saddr == sockaddr[0] or not sfamily == family:
                            raise PolicyError("HostStruct: '%s' resolves to multiple IP addresses" %
                                              hostname)
            if not foundFirst:
                raise PolicyError("HostStruct: '%s' did not resolve to one of the supported address family" %
                        hostname)
            self.name = hostname
            self.saddr = saddr
            self.family = sfamily
            self.binary = socket.inet_pton(family, saddr)
            return
        except Exception as e:
            raise PolicyError("HostStruct: '%s' failed to resolve: '%s'" %
                              (hostname, e))

    def __str__(self):
        return self.name

    def __repr__(self):
        return self.__str__()

    def dump(self):
        return ("(%s, %s, %s, %s)" %
                (self.name,
                 self.saddr,
                 "AF_INET" if self.family == socket.AF_INET else "AF_INET6",
                 binascii.hexlify(self.binary)))

#
#
class HostAddr(object):
    """
    Provide HostIP address ranges and comparison functions.
    A HostIP may be:
    - single address:      10.10.1.1
    - a pair of addresses: 10.10.0.0,10.10.255.255
    - a wildcard:          *
    Only IPv4 and IPv6 are supported.
    - No unix sockets.
    HostIP names must resolve to a single IP address.
    Address pairs define a range.
    - The second address must be numerically larger than the first address.
    - The addresses must be of the same address 'family', IPv4 or IPv6.
    The wildcard '*' matches all address IPv4 or IPv6.
    IPv6 support is conditional based on underlying OS network options.
    Raises a PolicyError on validation error in constructor.
    """

    def __init__(self, hostspec, separator=","):
        """
        Parse host spec into binary structures to use for comparisons.
        Validate the hostspec to enforce usage rules.
        """
        self.hoststructs = []

        if hostspec == "*":
            self.wildcard = True
        else:
            self.wildcard = False

            hosts = [x.strip() for x in hostspec.split(separator)]

            # hosts must contain one or two host specs
            if len(hosts) not in [1, 2]:
                raise PolicyError("hostspec must contain 1 or 2 host names")
            self.hoststructs.append(HostStruct(hosts[0]))
            if len(hosts) > 1:
                self.hoststructs.append(HostStruct(hosts[1]))
                if not self.hoststructs[0].family == self.hoststructs[1].family:
                    raise PolicyError("mixed IPv4 and IPv6 host specs in range not allowed")
                c0 = self.memcmp(self.hoststructs[0].binary, self.hoststructs[1].binary)
                if c0 > 0:
                    raise PolicyError("host specs in range must have lower numeric address first")

    def __str__(self):
        if self.wildcard:
            return "*"
        res = self.hoststructs[0].name
        if len(self.hoststructs) > 1:
            res += "," + self.hoststructs[1].name
        return res

    def __repr__(self):
        return self.__str__()

    def dump(self):
        if self.wildcard:
            return "(*)"
        res = "(" + self.hoststructs[0].dump()
        if len(self.hoststructs) > 1:
            res += "," + self.hoststructs[1].dump()
        res += ")"
        return res

    def memcmp(self, a, b):
        res = 0
        for i in range(0,len(a)):
            if a[i] > b[i]:
                res = 1
                break;
            elif a[i] < b[i]:
                res = -1
                break
        return res

    def match_bin(self, candidate):
        """
        Does the candidate hoststruct match the IP or range of IP addresses represented by this?
        @param[in] candidate the IP address to be tested
        @return candidate matches this or not
        """
        if self.wildcard:
            return True
        try:
            if not candidate.family == self.hoststructs[0].family:
                # sorry, wrong AF_INET family
                return False
            c0 = self.memcmp(candidate.binary, self.hoststructs[0].binary)
            if len(self.hoststructs) == 1:
                return c0 == 0
            c1 = self.memcmp(candidate.binary, self.hoststructs[1].binary)
            return c0 >= 0 and c1 <= 0
        except PolicyError:
            return False
        except Exception as e:
            assert isinstance(candidate, HostStruct), \
                ("Wrong type. Expected HostStruct but received %s" % candidate.__class__.__name__)
            return False

    def match_str(self, candidate):
        """
        Does the candidate string match the IP or range represented by this?
        @param[in] candidate the IP address to be tested
        @return candidate matches this or not
        """
        try:
            hoststruct = HostStruct(candidate)
        except PolicyError:
            return False
        return self.match_bin(hoststruct)

#
#
class PolicyAppConnectionMgr(object):
    """
    Track policy user/host connection limits and statistics for one app.
    # limits - set at creation and by update()
    max_total            : 20
    max_per_user         : 5
    max_per_host         : 10
    # statistics - maintained for the lifetime of corresponding application
    connections_approved : N
    connections_denied   : N
    # live state - maintained for the lifetime of corresponding application
    connections_active   : 5
    per_host_state : { 'host1' : [conn1, conn2, conn3],
                       'host2' : [conn4, conn5] }
    per_user_state : { 'user1' : [conn1, conn2, conn3],
                       'user2' : [conn4, conn5] }
    """
    def __init__(self, maxconn, maxconnperuser, maxconnperhost):
        """
        The object is constructed with the policy limits and zeroed counts.
        @param[in] maxconn maximum total concurrent connections
        @param[in] maxconnperuser maximum total conncurrent connections for each user
        @param[in] maxconnperuser maximum total conncurrent connections for each host
        """
        if maxconn < 0 or maxconnperuser < 0 or maxconnperhost < 0:
            raise PolicyError("PolicyAppConnectionMgr settings must be >= 0")
        self.max_total    = maxconn
        self.max_per_user = maxconnperuser
        self.max_per_host = maxconnperhost
        self.connections_approved = 0
        self.connections_denied   = 0
        self.connections_active   = 0
        self.per_host_state = {}
        self.per_user_state = {}

    def __str__(self):
        res = ("Connection Limits: total: %s, per user: %s, per host: %s\n" %
            (self.max_total, self.max_per_user, self.max_per_host))
        res += ("Connections Statistics: total approved: %s, total denied: %s" %
                (self.connections_approved, self.connections_denied))
        res += ("Connection State: total current: %s" % self.connections_active)
        res += ("User state: %s\n" % self.per_user_state)
        res += ("Host state: %s"   % self.per_host_state)
        return res

    def __repr__(self):
        return self.__str__()

    def update(self, maxconn, maxconnperuser, maxconnperhost):
        """
        Reset connection limits
        @param[in] maxconn maximum total concurrent connections
        @param[in] maxconnperuser maximum total conncurrent connections for each user
        @param[in] maxconnperuser maximum total conncurrent connections for each host
        """
        if maxconn < 0 or maxconnperuser < 0 or maxconnperhost < 0:
            raise PolicyError("PolicyAppConnectionMgr settings must be >= 0")
        self.max_total    = maxconn
        self.max_per_user = maxconnperuser
        self.max_per_host = maxconnperhost

    def can_connect(self, conn_id, user, host, diags):
        """
        Register a connection attempt.
        If all the connection limit rules pass then add the
        user/host to the connection tables.
        @param[in] conn_id unique ID for connection, usually IP:port
        @param[in] user authenticated user ID
        @param[in] host IP address of host
        @param[out] diags on failure holds 1, 2, or 3 error strings
        @return connection is allowed and tracked in state tables
        """
        n_user = 0
        if user in self.per_user_state:
            n_user = len(self.per_user_state[user])
        n_host = 0
        if host in self.per_host_state:
            n_host = len(self.per_host_state[host])

        allowbytotal = self.connections_active < self.max_total
        allowbyuser  = n_user < self.max_per_user
        allowbyhost  = n_host < self.max_per_host

        if allowbytotal and allowbyuser and allowbyhost:
            if not user in self.per_user_state:
                self.per_user_state[user] = []
            self.per_user_state[user].append(conn_id)
            if not host in self.per_host_state:
                self.per_host_state[host] = []
            self.per_host_state[host].append(conn_id)
            self.connections_active += 1
            self.connections_approved += 1
            return True
        else:
            if not allowbytotal:
                diags.append("Connection denied by application connection limit")
            if not allowbyuser:
                diags.append("Connection denied by application per user limit")
            if not allowbyhost:
                diags.append("Connection denied by application per host limit")
            self.connections_denied += 1
            return False

    def disconnect(self, conn_id, user, host):
        """
        Unregister a connection
        """
        assert(self.connections_active > 0)
        assert(user in self.per_user_state)
        assert(conn_id in self.per_user_state[user])
        assert(conn_id in self.per_host_state[host])
        self.connections_active -= 1
        self.per_user_state[user].remove(conn_id)
        self.per_host_state[host].remove(conn_id)

    def count_other_denial(self):
        """
        Record the statistic for a connection denied by some other process
        @return:
        """
        self.connections_denied += 1
<commit_msg>DISPATCH-1462: close socket used in IPv6 availability test<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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
#

from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function

import socket
import binascii

#
#
class PolicyError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return str(self.value)

def is_ipv6_enabled():
    """
    Returns true if IPV6 is enabled, false otherwise
    """
    ipv6_enabled = True
    try:
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.bind(('::1', 0))
        sock.close()
    except Exception as e:
        ipv6_enabled = False

    return ipv6_enabled

class HostStruct(object):
    """
    HostStruct represents a single, binary socket address from getaddrinfo
        - name     : name given to constructor; numeric IP or host name
        - saddr    : net name resolved by getaddrinfo; numeric IP
        - family   : saddr.family; int
        - binary   : saddr packed binary address; binary string
    """
    families = [socket.AF_INET]
    famnames = ["IPv4"]
    if is_ipv6_enabled():
        families.append(socket.AF_INET6)
        famnames.append("IPv6")

    def __init__(self, hostname):
        """
        Given a host name text string, return the socket info for it.
        @param[in] hostname host IP address to parse
        """
        try:
            res = socket.getaddrinfo(hostname, 0, socket.AF_UNSPEC, socket.SOCK_STREAM)
            if len(res) == 0:
                raise PolicyError("HostStruct: '%s' did not resolve to an IP address" % hostname)
            foundFirst = False
            saddr = ""
            sfamily = socket.AF_UNSPEC
            for i0 in range(0, len(res)):
                family, dum0, dum1, dum2, sockaddr = res[i0]
                if not foundFirst:
                    if family in self.families:
                        saddr = sockaddr[0]
                        sfamily = family
                        foundFirst = True
                else:
                    if family in self.families:
                        if not saddr == sockaddr[0] or not sfamily == family:
                            raise PolicyError("HostStruct: '%s' resolves to multiple IP addresses" %
                                              hostname)
            if not foundFirst:
                raise PolicyError("HostStruct: '%s' did not resolve to one of the supported address family" %
                        hostname)
            self.name = hostname
            self.saddr = saddr
            self.family = sfamily
            self.binary = socket.inet_pton(family, saddr)
            return
        except Exception as e:
            raise PolicyError("HostStruct: '%s' failed to resolve: '%s'" %
                              (hostname, e))

    def __str__(self):
        return self.name

    def __repr__(self):
        return self.__str__()

    def dump(self):
        return ("(%s, %s, %s, %s)" %
                (self.name,
                 self.saddr,
                 "AF_INET" if self.family == socket.AF_INET else "AF_INET6",
                 binascii.hexlify(self.binary)))

#
#
class HostAddr(object):
    """
    Provide HostIP address ranges and comparison functions.
    A HostIP may be:
    - single address:      10.10.1.1
    - a pair of addresses: 10.10.0.0,10.10.255.255
    - a wildcard:          *
    Only IPv4 and IPv6 are supported.
    - No unix sockets.
    HostIP names must resolve to a single IP address.
    Address pairs define a range.
    - The second address must be numerically larger than the first address.
    - The addresses must be of the same address 'family', IPv4 or IPv6.
    The wildcard '*' matches all address IPv4 or IPv6.
    IPv6 support is conditional based on underlying OS network options.
    Raises a PolicyError on validation error in constructor.
    """

    def __init__(self, hostspec, separator=","):
        """
        Parse host spec into binary structures to use for comparisons.
        Validate the hostspec to enforce usage rules.
        """
        self.hoststructs = []

        if hostspec == "*":
            self.wildcard = True
        else:
            self.wildcard = False

            hosts = [x.strip() for x in hostspec.split(separator)]

            # hosts must contain one or two host specs
            if len(hosts) not in [1, 2]:
                raise PolicyError("hostspec must contain 1 or 2 host names")
            self.hoststructs.append(HostStruct(hosts[0]))
            if len(hosts) > 1:
                self.hoststructs.append(HostStruct(hosts[1]))
                if not self.hoststructs[0].family == self.hoststructs[1].family:
                    raise PolicyError("mixed IPv4 and IPv6 host specs in range not allowed")
                c0 = self.memcmp(self.hoststructs[0].binary, self.hoststructs[1].binary)
                if c0 > 0:
                    raise PolicyError("host specs in range must have lower numeric address first")

    def __str__(self):
        if self.wildcard:
            return "*"
        res = self.hoststructs[0].name
        if len(self.hoststructs) > 1:
            res += "," + self.hoststructs[1].name
        return res

    def __repr__(self):
        return self.__str__()

    def dump(self):
        if self.wildcard:
            return "(*)"
        res = "(" + self.hoststructs[0].dump()
        if len(self.hoststructs) > 1:
            res += "," + self.hoststructs[1].dump()
        res += ")"
        return res

    def memcmp(self, a, b):
        res = 0
        for i in range(0,len(a)):
            if a[i] > b[i]:
                res = 1
                break;
            elif a[i] < b[i]:
                res = -1
                break
        return res

    def match_bin(self, candidate):
        """
        Does the candidate hoststruct match the IP or range of IP addresses represented by this?
        @param[in] candidate the IP address to be tested
        @return candidate matches this or not
        """
        if self.wildcard:
            return True
        try:
            if not candidate.family == self.hoststructs[0].family:
                # sorry, wrong AF_INET family
                return False
            c0 = self.memcmp(candidate.binary, self.hoststructs[0].binary)
            if len(self.hoststructs) == 1:
                return c0 == 0
            c1 = self.memcmp(candidate.binary, self.hoststructs[1].binary)
            return c0 >= 0 and c1 <= 0
        except PolicyError:
            return False
        except Exception as e:
            assert isinstance(candidate, HostStruct), \
                ("Wrong type. Expected HostStruct but received %s" % candidate.__class__.__name__)
            return False

    def match_str(self, candidate):
        """
        Does the candidate string match the IP or range represented by this?
        @param[in] candidate the IP address to be tested
        @return candidate matches this or not
        """
        try:
            hoststruct = HostStruct(candidate)
        except PolicyError:
            return False
        return self.match_bin(hoststruct)

#
#
class PolicyAppConnectionMgr(object):
    """
    Track policy user/host connection limits and statistics for one app.
    # limits - set at creation and by update()
    max_total            : 20
    max_per_user         : 5
    max_per_host         : 10
    # statistics - maintained for the lifetime of corresponding application
    connections_approved : N
    connections_denied   : N
    # live state - maintained for the lifetime of corresponding application
    connections_active   : 5
    per_host_state : { 'host1' : [conn1, conn2, conn3],
                       'host2' : [conn4, conn5] }
    per_user_state : { 'user1' : [conn1, conn2, conn3],
                       'user2' : [conn4, conn5] }
    """
    def __init__(self, maxconn, maxconnperuser, maxconnperhost):
        """
        The object is constructed with the policy limits and zeroed counts.
        @param[in] maxconn maximum total concurrent connections
        @param[in] maxconnperuser maximum total conncurrent connections for each user
        @param[in] maxconnperuser maximum total conncurrent connections for each host
        """
        if maxconn < 0 or maxconnperuser < 0 or maxconnperhost < 0:
            raise PolicyError("PolicyAppConnectionMgr settings must be >= 0")
        self.max_total    = maxconn
        self.max_per_user = maxconnperuser
        self.max_per_host = maxconnperhost
        self.connections_approved = 0
        self.connections_denied   = 0
        self.connections_active   = 0
        self.per_host_state = {}
        self.per_user_state = {}

    def __str__(self):
        res = ("Connection Limits: total: %s, per user: %s, per host: %s\n" %
            (self.max_total, self.max_per_user, self.max_per_host))
        res += ("Connections Statistics: total approved: %s, total denied: %s" %
                (self.connections_approved, self.connections_denied))
        res += ("Connection State: total current: %s" % self.connections_active)
        res += ("User state: %s\n" % self.per_user_state)
        res += ("Host state: %s"   % self.per_host_state)
        return res

    def __repr__(self):
        return self.__str__()

    def update(self, maxconn, maxconnperuser, maxconnperhost):
        """
        Reset connection limits
        @param[in] maxconn maximum total concurrent connections
        @param[in] maxconnperuser maximum total conncurrent connections for each user
        @param[in] maxconnperuser maximum total conncurrent connections for each host
        """
        if maxconn < 0 or maxconnperuser < 0 or maxconnperhost < 0:
            raise PolicyError("PolicyAppConnectionMgr settings must be >= 0")
        self.max_total    = maxconn
        self.max_per_user = maxconnperuser
        self.max_per_host = maxconnperhost

    def can_connect(self, conn_id, user, host, diags):
        """
        Register a connection attempt.
        If all the connection limit rules pass then add the
        user/host to the connection tables.
        @param[in] conn_id unique ID for connection, usually IP:port
        @param[in] user authenticated user ID
        @param[in] host IP address of host
        @param[out] diags on failure holds 1, 2, or 3 error strings
        @return connection is allowed and tracked in state tables
        """
        n_user = 0
        if user in self.per_user_state:
            n_user = len(self.per_user_state[user])
        n_host = 0
        if host in self.per_host_state:
            n_host = len(self.per_host_state[host])

        allowbytotal = self.connections_active < self.max_total
        allowbyuser  = n_user < self.max_per_user
        allowbyhost  = n_host < self.max_per_host

        if allowbytotal and allowbyuser and allowbyhost:
            if not user in self.per_user_state:
                self.per_user_state[user] = []
            self.per_user_state[user].append(conn_id)
            if not host in self.per_host_state:
                self.per_host_state[host] = []
            self.per_host_state[host].append(conn_id)
            self.connections_active += 1
            self.connections_approved += 1
            return True
        else:
            if not allowbytotal:
                diags.append("Connection denied by application connection limit")
            if not allowbyuser:
                diags.append("Connection denied by application per user limit")
            if not allowbyhost:
                diags.append("Connection denied by application per host limit")
            self.connections_denied += 1
            return False

    def disconnect(self, conn_id, user, host):
        """
        Unregister a connection
        """
        assert(self.connections_active > 0)
        assert(user in self.per_user_state)
        assert(conn_id in self.per_user_state[user])
        assert(conn_id in self.per_host_state[host])
        self.connections_active -= 1
        self.per_user_state[user].remove(conn_id)
        self.per_host_state[host].remove(conn_id)

    def count_other_denial(self):
        """
        Record the statistic for a connection denied by some other process
        @return:
        """
        self.connections_denied += 1
<|endoftext|>


<commit_before>'''
Unit tests for the robotutils.robotcomm module.

This is a port of the JUnit test for the corresponding Java
robotutils class, ConfigurationHelper
Author: JMJ
'''

import unittest
import random
import logging
import collections
import time

from . import robotcomm as rc
from .. import concurrent_helper as ch

logger = logging.getLogger(__name__) # pylint: disable=invalid-name

def trace(*args, **kwargs):
    """simply call debug"""
    logger.debug(args, kwargs)

def tracing():
    """ Whether to trace or not"""
    return True

TransportStats = collections.namedtuple('TransportStats', 'sends recvs forcedrops randomdrops')

#pylint: disable=invalid-name
class MockRemoteNode(rc.DatagramTransport.RemoteNode):
    """Test remote node implementation"""

    def __init__(self, transport, address):
        """Construct a node under {transport} with the specified {address}"""
        self._transport = transport
        self._address = address

    def send(self, msg) -> None:
        """Sends {msg} over transport"""
        self._transport.numsends.next()
        if self._force_drop(msg):
            if tracing():
                trace("TRANSPORT FORCEDROP:\n[%s]", msg)
            self._transport.numforcedrops.next()
            return # **************** EARLY RETURN *****

        if self._nonforce_drop():
            if tracing():
                trace("TRANSPORT: RANDDROP\n[%s]", msg)
            self._transport.numrandomdrops.next()
            return # **************** EARLY RETURN *****

        self._transport.handleSend(msg)

    def address(self) -> str:
        """Text representation of destination address"""
        return self._address

    def isrunning(self) -> bool:
        """Returns if the transport is able to send and receive packets"""
        return self._transport.isrunning()

    def _force_drop(self, msg) -> bool:
        """Whether or not to always drop this message"""
        return self._transport.ALWAYSDROP_TEXT in msg

    def _nonforce_drop(self) -> bool:
        """Whether randomly drop message"""
        if self._transport.closed:
            return True
        if self._transport.zeroFailures():
            return False
        return random.random() < self._transport.failurerate

    def __str__(self):
        return self._address

class MockTransport(rc.DatagramTransport): # pylint: disable=too-many-instance-attributes
    """Implements the mock transport used in these tests"""

    ALWAYSDROP_TEXT = "TRANSPORT-MUST-DROP"

    def __init__(self, local_address):
        """Initialize MockTransport. {local_address}" can be any string."""
        self.loopbacknode = MockRemoteNode(self, local_address)
        self.recvqueue = ch.ConcurrentDeque()
        self.transport_timer = None
        self.numsends = ch.AtomicNumber(0)
        self.numrecvs = ch.AtomicNumber(0)
        self.numforcedrops = ch.AtomicNumber(0)
        self.numrandomdrops = ch.AtomicNumber(0)
        self.closed = False
        self.failurerate = 0
        self.maxdelay = 0
        self.client_recv = None # set in start_listening
        self.scheduler = ch.EventScheduler()

    #
    # DatagramTransport methods
    #
    def start_listening(self, handler) -> None:
        """Listen should only be called once."""
        assert not self.client_recv
        self.client_recv = handler
        self.scheduler.start()

    def stop_listening(self) -> None:
        """Stop previously-started listening sessions."""
        assert self.client_recv
        self.client_recv = None

    def new_remotenode(self, address) -> rc.DatagramTransport.RemoteNode:
        """COnstruct and return a new remote node"""
        return MockRemoteNode(self, address)

    def close(self) -> None:
        """Close the transport"""
        self.scheduler.cancel_all()
        # print("Stopping scheduler...")
        self.scheduler.stop(block=True)
        # print("Scheduler stopped")
        self.closed = True
        if self.transport_timer:
            self.transport_timer.cancel()

    #
    # Methods to be used by tests
    #

    def setTransportCharacteristics(self, failurerate, maxdelay) -> None:
        """Changes transport characteristics: sends datagrams with {failurerate} failure
        rate, and {maxdelay} max delay per packet."""
        assert 0 <= failurerate <= 1
        assert maxdelay >= 0
        self.failurerate = failurerate
        self.maxdelay = maxdelay

    def healthy(self) -> bool:
        """Returns if the underlying transport is healthy, i.e., can still send/recv messages."""
        return self.scheduler.healthy()

    def zeroFailures(self) -> bool:
        """{failurerate} "close enough" to 0 is considered to be zero-failures."""
        return abs(self.failurerate) < 1e-7

    def noDelays(self) -> bool:
        """Returns True IFF the transport should never delay"""
        return self.maxdelay == 0

    def getStats(self) -> TransportStats:
        """Return transport statistics"""
        return TransportStats(
            sends=self.numsends.value(),
            recvs=self.numrecvs.value(),
            forcedrops=self.numforcedrops.value(),
            randomdrops=self.numrandomdrops.value())

    def getMaxDelay(self) -> float:
        """Return the max statistical delay between packets"""
        return self.maxdelay

    def getFailureRate(self) -> float:
        """Return the statistical packet failure rate"""
        return self.failurerate

    def handleSend(self, msg) -> None:
        """Internally handle a send request"""
        if self.closed:
            self.numforcedrops.next()
            return
        if self.noDelays():
            self._receiveData(msg)
        else:
            delay = random.random() * self.maxdelay
            def timertask():
                if not self.closed:
                    # Delayed send
                    self._receiveData(msg)
                else:
                    self.numforcedrops.next()

            self.scheduler.schedule(delay, timertask)

    #
    # Private methods
    #

    def _receiveData(self, s) -> None:
        """Internally handle a received message"""

        if self.client_recv:
            self.numrecvs.next()
            if tracing():
                trace("TRANSPORT RECV:\n[{}]\n", s)
            self.client_recv(s, self.loopbacknode)
        else:
            self.numforcedrops.next()

class TestMockTransport(unittest.TestCase):
    """Container for mock transport unit tests"""

    LOCAL_ADDRESS = "loopback"

    @staticmethod
    def cointoss():
        """Rreturns True with 0.5 probabilyt"""
        return random.random() < 0.5

    def test_mock_transport_without_random_failures(self): # pylint: disable=too-many-locals
        """Sends a bunch of messages and verifies each one is received exactly once."""
        msgcount = 10000
        expected_msgcount = 0
        failurerate = 0
        maxdelay = 1

        def genmsg(x):
            nonlocal expected_msgcount
            msg = 'msg-' + str(x)
            if self.cointoss():
                expected_msgcount += 1
                return msg
            # force drop
            return msg + MockTransport.ALWAYSDROP_TEXT

        messages = {genmsg(i): ch.AtomicNumber(0) for i in range(msgcount)}
        recvcount = ch.AtomicNumber(0)
        errcount = ch.AtomicNumber(0) # count of unexpected messages recvd

        def receivemsg(msg, node):
            # print("Received message [{}] from {}".format(msg, node))
            if node.address() != self.LOCAL_ADDRESS:
                errcount.next()
            recvcount.next()
            counter = messages.get(msg)
            if counter:
                counter.next()
            else:
                errcount.next() # unexpected message

        transport = MockTransport(self.LOCAL_ADDRESS)
        transport.setTransportCharacteristics(failurerate, maxdelay)
        transport.start_listening(receivemsg)
        node = transport.new_remotenode(self.LOCAL_ADDRESS)
        for msg in messages:
            node.send(msg) # We're sending the keys
        # print("Waiting... for {} messages".format(expected_msgcount))
        while recvcount.value() < expected_msgcount and transport.healthy():
            time.sleep(0.1)
        # print("Received {} messages".format(msgcount))

        # Need to close before any test failure assertions, otherwise
        # scheduler thread will never exit
        transport.close()

        self.assertEqual(errcount.value(), 0)
        self.assertEqual(recvcount.value(), expected_msgcount)

        # Check that expected messages are received exactly once, and
        # unexpected messages are not received.
        for msg, counter in messages.items():
            expected = 1
            if  MockTransport.ALWAYSDROP_TEXT in msg:
                expected = 0
            self.assertEqual(counter.value(), expected, "msg count for:" + msg)

        print("transport_simple: received {}/{} messages".format(recvcount, msgcount))

    def test_mock_transport_with_random_failures(self):
        """Sends a bunch of messages with random transport failures."""
        msgcount = 10000
        failurerate = 0.3
        maxdelay = 1

        recvcount = ch.AtomicNumber(0)
        messages = {'msg-' + str(i) for i in range(msgcount)}
        transport = MockTransport(self.LOCAL_ADDRESS)
        transport.setTransportCharacteristics(failurerate, maxdelay)

        # def receivemsg(msg, node):
        #     recvcount.next()
        transport.start_listening(lambda msg, node: recvcount.next())

        node = transport.new_remotenode("loopback")
        for msg in messages:
            node.send(msg) # We're sending the keys
        time.sleep(maxdelay)

        # Need to close before any test failure assertions, otherwise
        # scheduler thread will never exit
        transport.close()

        # Actual failure rate must match expected to within 0.1
        self.assertAlmostEqual(recvcount.value()/msgcount, 1-failurerate, 1)
        print("transport_simple: received {}/{} messages".format(recvcount, msgcount))

# Keeps track of a single test message
TestMessageRecord = collections.namedtuple('TestMessageRecord', 'id alwaysDrop msgType msgBody')

# Keeps track of a command and expected response
TestCommandRecord = collections.namedtuple('TestCommandRecord', 'cmdRecord respRecord rt timeout')

def new_message_record(id_, alwaysdrop):
    """Creates and returns a new test message record"""

    def rand32():
        return random.randint(0, 1<<32)

    def random_body():
        """First line in message body is the internal 'id'. Remainder is random junk."""
        extra = random.randint(1, 9)
        sequence = (hex(random.randint(0, 1<<32)) for _ in range(extra))
        return hex(id_) + '\n' + '\n'.join(sequence)

    def random_type():
        x = hex(rand32())
        return x[:random.randint(0, len(x))] # can be the empty string

    body = MockTransport.ALWAYSDROP_TEXT if alwaysdrop else random_body()
    return TestMessageRecord(id=id_, alwaysDrop=alwaysdrop,
                             msgType=random_type(), msgBody=body)

# class StressTester:
#     """Test engine that instantiates and runs robotcomm stress tests with
#     various parameters"""

class TestRobotComm(unittest.TestCase):
    """Container for RobotComm unit tests"""

    @unittest.skip("Unimplemented")
    def test_message_basic(self):
        """Sends and receives a single message"""
        transport = MockTransport("localhost")
        rcomm = rc.RobotComm(transport)
        channel = rcomm.new_channel("testChannel")
        channel.bind_to_remote_node("localhost")
        rcomm.start_listening()

        testMessage = "test message"
        channel.start_receiving_messages()
        channel.send_message("MYTYPE", testMessage)

        rm = None
        while not rm:
            rm = channel.poll_received_message()
            time.sleep(0.01)

        for stats in rcomm.get_channel_statistics():
            print(stats)

        channel.stop_receiving_messages()
        rcomm.stop_listening()
        rcomm.close()
        self.assertTrue(rm)
        self.assertEqual(testMessage, rm.message())
<commit_msg>Fix syntax errors in test_robotcomm.py<commit_after>'''
Unit tests for the robotutils.robotcomm module.

This is a port of the JUnit test for the corresponding Java
robotutils class, ConfigurationHelper
Author: JMJ
'''

import unittest
import random
import logging
import collections
import time

from . import robotcomm as rc
from .. import concurrent_helper as ch
from .common import DatagramTransport

logger = logging.getLogger(__name__) # pylint: disable=invalid-name

def trace(*args, **kwargs):
    """simply call debug"""
    logger.debug(args, kwargs)

def tracing():
    """ Whether to trace or not"""
    return True

TransportStats = collections.namedtuple('TransportStats', 'sends recvs forcedrops randomdrops')

#pylint: disable=invalid-name
class MockRemoteNode(DatagramTransport.RemoteNode):
    """Test remote node implementation"""

    def __init__(self, transport, address):
        """Construct a node under {transport} with the specified {address}"""
        self._transport = transport
        self._address = address

    def send(self, msg) -> None:
        """Sends {msg} over transport"""
        self._transport.numsends.next()
        if self._force_drop(msg):
            if tracing():
                trace("TRANSPORT FORCEDROP:\n[%s]", msg)
            self._transport.numforcedrops.next()
            return # **************** EARLY RETURN *****

        if self._nonforce_drop():
            if tracing():
                trace("TRANSPORT: RANDDROP\n[%s]", msg)
            self._transport.numrandomdrops.next()
            return # **************** EARLY RETURN *****

        self._transport.handleSend(msg)

    def address(self) -> str:
        """Text representation of destination address"""
        return self._address

    def isrunning(self) -> bool:
        """Returns if the transport is able to send and receive packets"""
        return self._transport.isrunning()

    def _force_drop(self, msg) -> bool:
        """Whether or not to always drop this message"""
        return self._transport.ALWAYSDROP_TEXT in msg

    def _nonforce_drop(self) -> bool:
        """Whether randomly drop message"""
        if self._transport.closed:
            return True
        if self._transport.zeroFailures():
            return False
        return random.random() < self._transport.failurerate

    def __str__(self):
        return self._address

class MockTransport(DatagramTransport): # pylint: disable=too-many-instance-attributes
    """Implements the mock transport used in these tests"""

    ALWAYSDROP_TEXT = "TRANSPORT-MUST-DROP"

    def __init__(self, local_address):
        """Initialize MockTransport. {local_address}" can be any string."""
        self.loopbacknode = MockRemoteNode(self, local_address)
        self.recvqueue = ch.ConcurrentDeque()
        self.transport_timer = None
        self.numsends = ch.AtomicNumber(0)
        self.numrecvs = ch.AtomicNumber(0)
        self.numforcedrops = ch.AtomicNumber(0)
        self.numrandomdrops = ch.AtomicNumber(0)
        self.closed = False
        self.failurerate = 0
        self.maxdelay = 0
        self.client_recv = None # set in start_listening
        self.scheduler = ch.EventScheduler()

    #
    # DatagramTransport methods
    #
    def start_listening(self, handler) -> None:
        """Listen should only be called once."""
        assert not self.client_recv
        self.client_recv = handler
        self.scheduler.start()

    def stop_listening(self) -> None:
        """Stop previously-started listening sessions."""
        assert self.client_recv
        self.client_recv = None

    def new_remotenode(self, address) -> DatagramTransport.RemoteNode:
        """COnstruct and return a new remote node"""
        return MockRemoteNode(self, address)

    def close(self) -> None:
        """Close the transport"""
        self.scheduler.cancel_all()
        # print("Stopping scheduler...")
        self.scheduler.stop(block=True)
        # print("Scheduler stopped")
        self.closed = True
        if self.transport_timer:
            self.transport_timer.cancel()

    #
    # Methods to be used by tests
    #

    def setTransportCharacteristics(self, failurerate, maxdelay) -> None:
        """Changes transport characteristics: sends datagrams with {failurerate} failure
        rate, and {maxdelay} max delay per packet."""
        assert 0 <= failurerate <= 1
        assert maxdelay >= 0
        self.failurerate = failurerate
        self.maxdelay = maxdelay

    def healthy(self) -> bool:
        """Returns if the underlying transport is healthy, i.e., can still send/recv messages."""
        return self.scheduler.healthy()

    def zeroFailures(self) -> bool:
        """{failurerate} "close enough" to 0 is considered to be zero-failures."""
        return abs(self.failurerate) < 1e-7

    def noDelays(self) -> bool:
        """Returns True IFF the transport should never delay"""
        return self.maxdelay == 0

    def getStats(self) -> TransportStats:
        """Return transport statistics"""
        return TransportStats(
            sends=self.numsends.value(),
            recvs=self.numrecvs.value(),
            forcedrops=self.numforcedrops.value(),
            randomdrops=self.numrandomdrops.value())

    def getMaxDelay(self) -> float:
        """Return the max statistical delay between packets"""
        return self.maxdelay

    def getFailureRate(self) -> float:
        """Return the statistical packet failure rate"""
        return self.failurerate

    def handleSend(self, msg) -> None:
        """Internally handle a send request"""
        if self.closed:
            self.numforcedrops.next()
            return
        if self.noDelays():
            self._receiveData(msg)
        else:
            delay = random.random() * self.maxdelay
            def timertask():
                if not self.closed:
                    # Delayed send
                    self._receiveData(msg)
                else:
                    self.numforcedrops.next()

            self.scheduler.schedule(delay, timertask)

    #
    # Private methods
    #

    def _receiveData(self, s) -> None:
        """Internally handle a received message"""

        if self.client_recv:
            self.numrecvs.next()
            if tracing():
                trace("TRANSPORT RECV:\n[{}]\n", s)
            self.client_recv(s, self.loopbacknode)
        else:
            self.numforcedrops.next()

class TestMockTransport(unittest.TestCase):
    """Container for mock transport unit tests"""

    LOCAL_ADDRESS = "loopback"

    @staticmethod
    def cointoss():
        """Rreturns True with 0.5 probabilyt"""
        return random.random() < 0.5

    def test_mock_transport_without_random_failures(self): # pylint: disable=too-many-locals
        """Sends a bunch of messages and verifies each one is received exactly once."""
        msgcount = 10000
        expected_msgcount = 0
        failurerate = 0
        maxdelay = 1

        def genmsg(x):
            nonlocal expected_msgcount
            msg = 'msg-' + str(x)
            if self.cointoss():
                expected_msgcount += 1
                return msg
            # force drop
            return msg + MockTransport.ALWAYSDROP_TEXT

        messages = {genmsg(i): ch.AtomicNumber(0) for i in range(msgcount)}
        recvcount = ch.AtomicNumber(0)
        errcount = ch.AtomicNumber(0) # count of unexpected messages recvd

        def receivemsg(msg, node):
            # print("Received message [{}] from {}".format(msg, node))
            if node.address() != self.LOCAL_ADDRESS:
                errcount.next()
            recvcount.next()
            counter = messages.get(msg)
            if counter:
                counter.next()
            else:
                errcount.next() # unexpected message

        transport = MockTransport(self.LOCAL_ADDRESS)
        transport.setTransportCharacteristics(failurerate, maxdelay)
        transport.start_listening(receivemsg)
        node = transport.new_remotenode(self.LOCAL_ADDRESS)
        for msg in messages:
            node.send(msg) # We're sending the keys
        # print("Waiting... for {} messages".format(expected_msgcount))
        while recvcount.value() < expected_msgcount and transport.healthy():
            time.sleep(0.1)
        # print("Received {} messages".format(msgcount))

        # Need to close before any test failure assertions, otherwise
        # scheduler thread will never exit
        transport.close()

        self.assertEqual(errcount.value(), 0)
        self.assertEqual(recvcount.value(), expected_msgcount)

        # Check that expected messages are received exactly once, and
        # unexpected messages are not received.
        for msg, counter in messages.items():
            expected = 1
            if  MockTransport.ALWAYSDROP_TEXT in msg:
                expected = 0
            self.assertEqual(counter.value(), expected, "msg count for:" + msg)

        print("transport_simple: received {}/{} messages".format(recvcount, msgcount))

    def test_mock_transport_with_random_failures(self):
        """Sends a bunch of messages with random transport failures."""
        msgcount = 10000
        failurerate = 0.3
        maxdelay = 1

        recvcount = ch.AtomicNumber(0)
        messages = {'msg-' + str(i) for i in range(msgcount)}
        transport = MockTransport(self.LOCAL_ADDRESS)
        transport.setTransportCharacteristics(failurerate, maxdelay)

        # def receivemsg(msg, node):
        #     recvcount.next()
        transport.start_listening(lambda msg, node: recvcount.next())

        node = transport.new_remotenode("loopback")
        for msg in messages:
            node.send(msg) # We're sending the keys
        time.sleep(maxdelay)

        # Need to close before any test failure assertions, otherwise
        # scheduler thread will never exit
        transport.close()

        # Actual failure rate must match expected to within 0.1
        self.assertAlmostEqual(recvcount.value()/msgcount, 1-failurerate, 1)
        print("transport_simple: received {}/{} messages".format(recvcount, msgcount))

# Keeps track of a single test message
TestMessageRecord = collections.namedtuple('TestMessageRecord', 'id alwaysDrop msgType msgBody')

# Keeps track of a command and expected response
TestCommandRecord = collections.namedtuple('TestCommandRecord', 'cmdRecord respRecord rt timeout')

def new_message_record(id_, alwaysdrop):
    """Creates and returns a new test message record"""

    def rand32():
        return random.randint(0, 1<<32)

    def random_body():
        """First line in message body is the internal 'id'. Remainder is random junk."""
        extra = random.randint(1, 9)
        sequence = (hex(random.randint(0, 1<<32)) for _ in range(extra))
        return hex(id_) + '\n' + '\n'.join(sequence)

    def random_type():
        x = hex(rand32())
        return x[:random.randint(0, len(x))] # can be the empty string

    body = MockTransport.ALWAYSDROP_TEXT if alwaysdrop else random_body()
    return TestMessageRecord(id=id_, alwaysDrop=alwaysdrop,
                             msgType=random_type(), msgBody=body)

# class StressTester:
#     """Test engine that instantiates and runs robotcomm stress tests with
#     various parameters"""

class TestRobotComm(unittest.TestCase):
    """Container for RobotComm unit tests"""

    @unittest.skip("Unimplemented")
    def test_message_basic(self):
        """Sends and receives a single message"""
        transport = MockTransport("localhost")
        rcomm = rc.RobotComm(transport)
        channel = rcomm.new_channel("testChannel")
        channel.bind_to_remote_node("localhost")
        rcomm.start_listening()

        testMessage = "test message"
        channel.start_receiving_messages()
        channel.send_message("MYTYPE", testMessage)

        rm = None
        while not rm:
            rm = channel.poll_received_message()
            time.sleep(0.01)

        for stats in rcomm.get_channel_statistics():
            print(stats)

        channel.stop_receiving_messages()
        rcomm.stop_listening()
        rcomm.close()
        self.assertTrue(rm)
        self.assertEqual(testMessage, rm.message())
<|endoftext|>


<commit_before><commit_msg>FEM: Python base VP, better check in get icon<commit_after><|endoftext|>


<commit_before>import pytest
import pandas as pd
import duckdb
from cardpay_reward_programs.rules import safe_ownership, staking
import math

token_holder_table = "_TOKEN_HOLDER"
safe_ownership_table = "_SAFE_OWNERSHIP"

def create_rule(
    monkeypatch, fake_data_token_holder, fake_data_safe_owner, core_config_overrides={}, user_config_overrides={}
):
    core_config = {
        "payment_cycle_length": 30,
        "start_block": 0,
        "end_block": 30,
        "subgraph_config_locations": {
            "prepaid_card_payment": "s3://cardpay-staging-partitioned-graph-data/data/prepaid_card_payments/0.0.3/",
            "spend_accumulation": "s3://tall-data-dev/paulin/spend_accumulation/0.0.1/",
            "safe_owner": "s3://tall-data-dev/paulin/safe_owner/0.0.1/"
        }
    }

    core_config.update(core_config_overrides)
    user_config = {
        "token": "card",
        "duration": 30,
        "interest_rate_monthly": 0.06
    }

    user_config.update(user_config_overrides)
    con = duckdb.connect(database=":memory:", read_only=False)
    con.execute(f"""create table {token_holder_table} as select * from fake_data_token_holder""")
    con.execute(f"""create table {safe_ownership_table} as select * from fake_data_safe_owner""")

    def table_query(self, config_name, table_name, min_partition: int, max_partition: int):

        if table_name == "token_holder":
            return token_holder_table
        elif table_name == "safe_owner":
            return safe_ownership_table
        else:
            raise NameError

    def run_query(self, table_query, vars, aux_table_query):
        con.execute(self.sql(token_holder_table, safe_ownership_table), vars)
        return con.fetchdf()

    monkeypatch.setattr(staking.Staking, "_get_table_query", table_query)
    monkeypatch.setattr(staking.Staking, "run_query", run_query)

    rule = staking.Staking(core_config, user_config)
    return rule

def get_amount(result, payee, idx):
    idx = result.index[result["payee"] == payee].tolist()[0]
    return result.loc[idx, "amount"]

def test_correctly_compound_interest(monkeypatch):
    fake_data_token_holder = pd.DataFrame(
        [
            {
                "_block_number": 0,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1000
            },
            {
                "_block_number": 10,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1500
            },
            {
                "_block_number": 15,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1200
            },
            {
                "_block_number": 20,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1700
            },
            {
                "_block_number": 25,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1900
            },
            {
                "_block_number": 0,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1500
            },
            {
                "_block_number": 0,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1500
            },
            {
                "_block_number": 15,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1200
            },
            {
                "_block_number": 25,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1000
            },
            {
                "_block_number": 0,
                "id": 3,
                "token": "card",
                "safe": "safe3",
                "balance_uint64": 1500
            },
        ]
    )

    fake_data_safe_owner = pd.DataFrame(
        [
            {
                "safe":"safe1",
                "owner":"owner1"
            },
            {
                "safe": "safe2",
                "owner": "owner2"
            },
            {
                "safe": "safe3",
                "owner": "owner3"
            }
        ]
    )

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(30, "0x0")
    assert len(result) == 3

def test_correct_calc_rewards(monkeypatch):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1000
        },
        {
            "_block_number": 10,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1500
        },
        {
            "_block_number": 20,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1750
        },
        {
            "_block_number": 25,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1500
        },
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1000
        },
        {
            "_block_number": 10,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1500
        },
        {
            "_block_number": 15,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1200
        },
        {
            "_block_number": 20,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1700
        },
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe3",
            "balance_uint64": 1000
        }

    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        },
        {
            "safe":"safe2",
            "owner":"owner2"
        },
        {
            "safe":"safe3",
            "owner":"owner3"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(30, "0x0")
    print(get_amount(result, "owner1", 0))
    assert math.ceil(get_amount(result, "owner1", 0)) == math.ceil(82.26866088)
    assert math.ceil(get_amount(result, "owner2", 1)) == math.ceil(80.74)
    assert math.ceil(get_amount(result, "owner3", 2)) == math.ceil(60)

def test_correctly_manages_first_deposit_in_cycle(monkeypatch):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1000
        },
        {
            "_block_number": 30,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 2000
        }
    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(60, "0x0")

    assert pytest.approx(get_amount(result, "owner1", 0)) == 120

def test_correct_calc_rewards_in_cycle(monkeypatch):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1000
        }
    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(30, "0x0")

    assert pytest.approx(get_amount(result, "owner1", 0)) == 60

def max_balance(monkeypatch):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 18446744073709551615
        }
    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )

    result = rule.run(30, "0x0")
    expected_value = 1.1068 * (10**18)
    assert get_amount(result, "owner1", 0) - expected_value < 4619683561473
<commit_msg>Issue with hypothesis positional parameters<commit_after>import pytest
import pandas as pd
import duckdb
import math
import hypothesis.strategies as st
import hypothesis.strategies as tuples
from cardpay_reward_programs.rules import safe_ownership, staking
from hypothesis import given
from hypothesis.extra.pandas import data_frames, columns

token_holder_table = "_TOKEN_HOLDER"
safe_ownership_table = "_SAFE_OWNERSHIP"
token_holder_columns = ["_block_number", "token", "safe", "balance_uint64"]
token_holder_df = data_frames(columns = columns(token_holder_columns, dtype = str),
                                rows = st.tuples(
                                    st.integers(min_value = 1000000, max_value = 18446744073709551615),
                                    st.from_regex("card\-[0-9]", fullmatch = True),
                                    st.from_regex("safe\-[0-9]", fullmatch = True),
                                    st.floats(min_value = 22995968, max_value = 24032768)
                                )
)
safe_owner_columns = ["safe", "owner"]
safe_owner_df = data_frames(columns = columns(safe_owner_columns, dtype = str),
                            rows = st.tuples(
                                st.from_regex("owner\-[0-9]", fullmatch = True),
                                st.from_regex("safe\-[0-9]", fullmatch = True),
                            )

                )

def create_rule(
    monkeypatch, fake_data_token_holder, fake_data_safe_owner, core_config_overrides={}, user_config_overrides={}
):
    core_config = {
        "payment_cycle_length": 30,
        "start_block": 0,
        "end_block": 30,
        "subgraph_config_locations": {
            "prepaid_card_payment": "s3://cardpay-staging-partitioned-graph-data/data/prepaid_card_payments/0.0.3/",
            "spend_accumulation": "s3://tall-data-dev/paulin/spend_accumulation/0.0.1/",
            "safe_owner": "s3://tall-data-dev/paulin/safe_owner/0.0.1/"
        }
    }

    core_config.update(core_config_overrides)
    user_config = {
        "token": "card",
        "duration": 30,
        "interest_rate_monthly": 0.06
    }

    user_config.update(user_config_overrides)
    con = duckdb.connect(database=":memory:", read_only=False)
    con.execute(f"""create table {token_holder_table} as select * from fake_data_token_holder""")
    con.execute(f"""create table {safe_ownership_table} as select * from fake_data_safe_owner""")

    def table_query(self, config_name, table_name, min_partition: int, max_partition: int):

        if table_name == "token_holder":
            return token_holder_table
        elif table_name == "safe_owner":
            return safe_ownership_table
        else:
            raise NameError

    def run_query(self, table_query, vars, aux_table_query):
        con.execute(self.sql(token_holder_table, safe_ownership_table), vars)
        return con.fetchdf()

    monkeypatch.setattr(staking.Staking, "_get_table_query", table_query)
    monkeypatch.setattr(staking.Staking, "run_query", run_query)

    rule = staking.Staking(core_config, user_config)
    return rule

def get_amount(result, payee, idx):
    idx = result.index[result["payee"] == payee].tolist()[0]
    return result.loc[idx, "amount"]

def test_correctly_compound_interest(monkeypatch):
    fake_data_token_holder = pd.DataFrame(
        [
            {
                "_block_number": 0,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1000
            },
            {
                "_block_number": 10,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1500
            },
            {
                "_block_number": 15,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1200
            },
            {
                "_block_number": 20,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1700
            },
            {
                "_block_number": 25,
                "id": 1,
                "token": "card",
                "safe": "safe1",
                "balance_uint64": 1900
            },
            {
                "_block_number": 0,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1500
            },
            {
                "_block_number": 0,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1500
            },
            {
                "_block_number": 15,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1200
            },
            {
                "_block_number": 25,
                "id": 2,
                "token": "card",
                "safe": "safe2",
                "balance_uint64": 1000
            },
            {
                "_block_number": 0,
                "id": 3,
                "token": "card",
                "safe": "safe3",
                "balance_uint64": 1500
            },
        ]
    )

    fake_data_safe_owner = pd.DataFrame(
        [
            {
                "safe":"safe1",
                "owner":"owner1"
            },
            {
                "safe": "safe2",
                "owner": "owner2"
            },
            {
                "safe": "safe3",
                "owner": "owner3"
            }
        ]
    )

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(30, "0x0")
    assert len(result) == 3

def test_correct_calc_rewards(monkeypatch):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1000
        },
        {
            "_block_number": 10,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1500
        },
        {
            "_block_number": 20,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1750
        },
        {
            "_block_number": 25,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1500
        },
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1000
        },
        {
            "_block_number": 10,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1500
        },
        {
            "_block_number": 15,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1200
        },
        {
            "_block_number": 20,
            "token": "card",
            "safe": "safe2",
            "balance_uint64": 1700
        },
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe3",
            "balance_uint64": 1000
        }

    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        },
        {
            "safe":"safe2",
            "owner":"owner2"
        },
        {
            "safe":"safe3",
            "owner":"owner3"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(30, "0x0")
    print(get_amount(result, "owner1", 0))
    assert math.ceil(get_amount(result, "owner1", 0)) == math.ceil(82.26866088)
    assert math.ceil(get_amount(result, "owner2", 1)) == math.ceil(80.74)
    assert math.ceil(get_amount(result, "owner3", 2)) == math.ceil(60)

def test_correctly_manages_first_deposit_in_cycle(monkeypatch):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1000
        },
        {
            "_block_number": 30,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 2000
        }
    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(60, "0x0")

    assert pytest.approx(get_amount(result, "owner1", 0)) == 120

def test_correct_calc_rewards_in_cycle(monkeypatch, i = 0, j = 0):
    fake_data_token_holder = pd.DataFrame([
        {
            "_block_number": 0,
            "token": "card",
            "safe": "safe1",
            "balance_uint64": 1000
        }
    ])

    fake_data_safe_owner = pd.DataFrame([
        {
            "safe":"safe1",
            "owner":"owner1"
        }
    ])

    rule = create_rule(
        monkeypatch, fake_data_token_holder, fake_data_safe_owner
    )
    result = rule.run(30, "0x0")

    assert pytest.approx(get_amount(result, "owner1", 0)) == 60

@given(token_holder_df, safe_owner_df)
def test_hypothesis(token_holder_df, safe_owner_df):
    assert len(token_holder_df) > -1

@given(st.integers(), st.integers())
def test_b(x, y, monkeypatch):
    print(x, y)
    assert 0 == 0

<|endoftext|>