<commit_before>from functools import partial
import signal
import multiprocessing
import asyncio
import time
import sys
import os

from molotov.util import log, stream_log
from molotov.session import LoggedClientSession as Session
from molotov.result import LiveResults, ClosedError
from molotov.api import get_fixture, pick_scenario, get_scenario
from molotov.stats import get_statsd_client
from molotov.ui import quit as quit_screen


_STOP = False
_REFRESH = .3


def _now():
    return int(time.time())


_results = LiveResults()


def get_live_results():
    return _results


async def consume(queue, numworkers, console=False, verbose=0):
    worker_stopped = 0
    while True and worker_stopped < numworkers:
        try:
            item = await queue.get()
        except RuntimeError:
            break
        if item == 'WORKER_STOPPED':
            worker_stopped += 1
        elif item == 'STOP':
            break

        elif isinstance(item, str):
            results = get_live_results()
            try:
                if item == '.':
                    results.incr_success()
                elif item == '-':
                    results.incr_failure()
                else:
                    if console and verbose > 0:
                        print(item)
            except ClosedError:
                break
        else:
            if console and verbose > 0:
                import traceback
                traceback.print_tb(item)


async def step(session, quiet, verbose, stream, scenario=None):
    """ single scenario call.

    When it returns 1, it works. -1 the script failed,
    0 the test is stopping or needs to stop.
    """
    if scenario:
        __, func, args_, kw = scenario
    else:
        func, args_, kw = pick_scenario()
    try:
        await func(session, *args_, **kw)
        await stream.put('.')
        return 1
    except asyncio.CancelledError:
        return 0
    except Exception as exc:
        await stream.put('-')
        if verbose > 0:
            await stream.put(repr(exc))
            await stream.put(sys.exc_info()[2])

    return -1


_HOWLONG = 0


async def worker(num, loop, results, args, stream, statsd):
    global _STOP
    quiet = args.quiet
    duration = args.duration
    verbose = args.verbose
    exception = args.exception
    if args.single_mode:
        single = get_scenario(args.single_mode)
    else:
        single = None
    count = 1
    start = _now()
    howlong = 0
    setup = get_fixture('setup')
    if setup is not None:
        try:
            options = await setup(num, args)
        except Exception as e:
            log(e)
            await stream.put('WORKER_STOPPED')
            return
        if options is None:
            options = {}
        elif not isinstance(options, dict):
            log('The setup function needs to return a dict')
            await stream.put('WORKER_STOPPED')
            return
    else:
        options = {}

    ssetup = get_fixture('setup_session')
    steardown = get_fixture('teardown_session')

    async with Session(loop, stream, verbose, statsd, **options) as session:
        session.args = args

        if ssetup is not None:
            await ssetup(num, session)

        while howlong < duration and not _STOP:
            if args.max_runs and count > args.max_runs:
                break
            howlong = _now() - start
            result = await step(session, quiet, verbose, stream, single)
            if result == 1:
                results['OK'] += 1
            elif result == -1:<commit_msg>added worker_id to the session<commit_after>from functools import partial
import signal
import multiprocessing
import asyncio
import time
import sys
import os

from molotov.util import log, stream_log
from molotov.session import LoggedClientSession as Session
from molotov.result import LiveResults, ClosedError
from molotov.api import get_fixture, pick_scenario, get_scenario
from molotov.stats import get_statsd_client
from molotov.ui import quit as quit_screen


_STOP = False
_REFRESH = .3


def _now():
    return int(time.time())


_results = LiveResults()


def get_live_results():
    return _results


async def consume(queue, numworkers, console=False, verbose=0):
    worker_stopped = 0
    while True and worker_stopped < numworkers:
        try:
            item = await queue.get()
        except RuntimeError:
            break
        if item == 'WORKER_STOPPED':
            worker_stopped += 1
        elif item == 'STOP':
            break

        elif isinstance(item, str):
            results = get_live_results()
            try:
                if item == '.':
                    results.incr_success()
                elif item == '-':
                    results.incr_failure()
                else:
                    if console and verbose > 0:
                        print(item)
            except ClosedError:
                break
        else:
            if console and verbose > 0:
                import traceback
                traceback.print_tb(item)


async def step(session, quiet, verbose, stream, scenario=None):
    """ single scenario call.

    When it returns 1, it works. -1 the script failed,
    0 the test is stopping or needs to stop.
    """
    if scenario:
        __, func, args_, kw = scenario
    else:
        func, args_, kw = pick_scenario()
    try:
        await func(session, *args_, **kw)
        await stream.put('.')
        return 1
    except asyncio.CancelledError:
        return 0
    except Exception as exc:
        await stream.put('-')
        if verbose > 0:
            await stream.put(repr(exc))
            await stream.put(sys.exc_info()[2])

    return -1


_HOWLONG = 0


async def worker(num, loop, results, args, stream, statsd):
    global _STOP
    quiet = args.quiet
    duration = args.duration
    verbose = args.verbose
    exception = args.exception
    if args.single_mode:
        single = get_scenario(args.single_mode)
    else:
        single = None
    count = 1
    start = _now()
    howlong = 0
    setup = get_fixture('setup')
    if setup is not None:
        try:
            options = await setup(num, args)
        except Exception as e:
            log(e)
            await stream.put('WORKER_STOPPED')
            return
        if options is None:
            options = {}
        elif not isinstance(options, dict):
            log('The setup function needs to return a dict')
            await stream.put('WORKER_STOPPED')
            return
    else:
        options = {}

    ssetup = get_fixture('setup_session')
    steardown = get_fixture('teardown_session')

    async with Session(loop, stream, verbose, statsd, **options) as session:
        session.args = args
        session.worker_id = num

        if ssetup is not None:
            await ssetup(num, session)

        while howlong < duration and not _STOP:
            if args.max_runs and count > args.max_runs:
                break
            howlong = _now() - start
            result = await step(session, quiet, verbose, stream, single)
            if result == 1:
                results['OK'] += 1
            elif result == -1:
{'fraction_non_alphanumeric': 0.08151342328557519, 'fraction_numerical': 0.005564056196967589, 'mean_word_length': 1.8748500599760096, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3167214', 'n_tokens_mistral': 2056, 'n_tokens_neox': 1969, 'n_words': 641}