# Path: kafka_influxdb/worker.py
# class Worker(object):
#     """
#     Implementation of worker class that handles Kafka and InfluxDB
#     connections and manages message encoding.
#     """
# 
#     def __init__(self, reader, encoder, writer, config):
#         """
#         Setup
#         """
#         self.config = config
#         self.reader = reader
#         self.encoder = encoder
#         self.writer = writer
#         self.buffer = []
# 
#         # Field for time measurement
#         self.start_time = None
#         self.last_flush_time = None
#         self.db_create_delay = 1
# 
#     def consume(self):
#         """
#         Run loop. Consume messages from reader, convert it to the
#         output format using encoder and write to output with writer
#         """
#         self.init_database()
# 
#         logging.info("Listening for messages on Kafka topic %s...",
#                      self.config.kafka_topic)
#         self.start_time = self.last_flush_time = time.time()
#         while True:
#             try:
#                 for index, raw_message in enumerate(self.reader.read(), 1):
#                     if raw_message:
#                         self.buffer.extend(self.encoder.encode(raw_message))
#                         if index % self.config.buffer_size == 0:
#                             self.flush()
#                     elif (self.config.buffer_timeout and len(self.buffer) > 0 and
#                           (time.time() - self.last_flush_time) >= self.config.buffer_timeout):
#                         logging.debug("Buffer timeout %ss. Flushing remaining %s messages from buffer.",
#                                       self.config.buffer_timeout, len(self.buffer))
#                         self.flush()
#             except EncoderError:
#                 logging.error("Encoder error. Trying to reconnect to %s:%s",
#                               self.config.kafka_host, self.config.kafka_port)
#                 logging.debug("Sleeping for %d ms before reconnect",
#                               self.config.reconnect_wait_time_ms)
#                 time.sleep(self.config.reconnect_wait_time_ms / 1000.0)
#             except KeyboardInterrupt:
#                 logging.info(
#                     "Shutdown. Flushing remaining messages from buffer.")
#                 self.flush()
#                 break
#             except SystemExit:
#                 break
# 
#     def init_database(self):
#         """
#         Initialize the InfluxDB database if it is not already there
#         """
#         try:
#             logging.info("Creating InfluxDB database if not exists: %s",
#                          self.config.influxdb_dbname)
#             self.writer.create_database(self.config.influxdb_dbname)
#         except ConnectionError as error:
#             logging.error(
#                 "Connection error while trying to create InfluxDB database: %s. Waiting for retry...", error)
#             time.sleep(self.db_create_delay)
#             self.init_database()
#         except (InfluxDBServerError, InfluxDBClientError) as error:
#             logging.warning(
#                 "Could not create InfluxDB database. Assuming it already exists: %s", error)
# 
#     def flush(self):
#         """
#         Flush values with writer
#         """
#         if not self.buffer:
#             # Don't do anything when buffer empty
#             return
#         try:
#             self.last_flush_time = time.time()
#             self.writer.write(self.buffer)
#             if self.config.statistics:
#                 self.show_statistics()
#         except (InfluxDBServerError, InfluxDBClientError) as influx_error:
#             logging.error("Error while writing to InfluxDB: %s", influx_error)
#         finally:
#             self.buffer = []
# 
#     def show_statistics(self):
#         """
#         Print performance metrics to stdout
#         """
#         delta = time.time() - self.start_time
#         msg_per_sec = self.config.buffer_size / delta
#         print("Flushing output buffer. {0:.2f} messages/s".format(msg_per_sec))
#         # Reset timer
#         self.start_time = time.time()
# 
#     def set_reader(self, reader):
#         self.reader = reader
# 
#     def get_reader(self):
#         return self.reader
# 
#     def set_writer(self, writer):
#         self.writer = writer
# 
#     def get_writer(self):
#         return self.writer
# 
#     def get_buffer(self):
#         return self.buffer
# 
#     def get_config(self):
#         return self.config
# 
# Path: kafka_influxdb/encoder/echo_encoder.py
# class Encoder(object):
#     def encode(msg):
# 
# Path: kafka_influxdb/tests/helpers/timeout.py
# def timeout(seconds=5):
#     def decorator(func):
#         def wrapper(*args, **kwargs):
#             process = Process(None, func, None, args, kwargs)
#             process.start()
#             process.join(seconds)
#             if process.is_alive():
#                 process.terminate()
#         return wraps(func)(wrapper)
#     return decorator


import unittest
import random
import time
from mock import Mock
from kafka_influxdb.worker import Worker
from kafka_influxdb.encoder import echo_encoder
from kafka_influxdb.tests.helpers.timeout import timeout
from requests.exceptions import ConnectionError
from mock import patch
        self.buffer_timeout = buffer_timeout
        self.kafka_topic = "test"
        self.influxdb_dbname = "mydb"
        self.statistics = False


class DummyReader(object):
    """
    A reader that yields dummy messages
    """

    def __init__(self, messages, num_messages):
        self.messages = messages
        self.num_messages = num_messages

    def read(self):
        for i in range(self.num_messages):
            yield random.choice(self.messages)
        # Simulate keyboard interrupt by user to stop consuming
        raise KeyboardInterrupt


class TimeoutReader(object):
    """
    A reader that writes half the messages then times out
    """

    def __init__(self, message, num_messages, timeout):
        self.message = message
        self.num_messages = num_messages

# Path: kafka_influxdb/worker.py
# class Worker(object):
#     """
#     Implementation of worker class that handles Kafka and InfluxDB
#     connections and manages message encoding.
#     """
# 
#     def __init__(self, reader, encoder, writer, config):
#         """
#         Setup
#         """
#         self.config = config
#         self.reader = reader
#         self.encoder = encoder
#         self.writer = writer
#         self.buffer = []
# 
#         # Field for time measurement
#         self.start_time = None
#         self.last_flush_time = None
#         self.db_create_delay = 1
# 
#     def consume(self):
#         """
#         Run loop. Consume messages from reader, convert it to the
#         output format using encoder and write to output with writer
#         """
#         self.init_database()
# 
#         logging.info("Listening for messages on Kafka topic %s...",
#                      self.config.kafka_topic)
#         self.start_time = self.last_flush_time = time.time()
#         while True:
#             try:
#                 for index, raw_message in enumerate(self.reader.read(), 1):
#                     if raw_message:
#                         self.buffer.extend(self.encoder.encode(raw_message))
#                         if index % self.config.buffer_size == 0:
#                             self.flush()
#                     elif (self.config.buffer_timeout and len(self.buffer) > 0 and
#                           (time.time() - self.last_flush_time) >= self.config.buffer_timeout):
#                         logging.debug("Buffer timeout %ss. Flushing remaining %s messages from buffer.",
#                                       self.config.buffer_timeout, len(self.buffer))
#                         self.flush()
#             except EncoderError:
#                 logging.error("Encoder error. Trying to reconnect to %s:%s",
#                               self.config.kafka_host, self.config.kafka_port)
#                 logging.debug("Sleeping for %d ms before reconnect",
#                               self.config.reconnect_wait_time_ms)
#                 time.sleep(self.config.reconnect_wait_time_ms / 1000.0)
#             except KeyboardInterrupt:
#                 logging.info(
#                     "Shutdown. Flushing remaining messages from buffer.")
#                 self.flush()
#                 break
#             except SystemExit:
#                 break
# 
#     def init_database(self):
#         """
#         Initialize the InfluxDB database if it is not already there
#         """
#         try:
#             logging.info("Creating InfluxDB database if not exists: %s",
#                          self.config.influxdb_dbname)
#             self.writer.create_database(self.config.influxdb_dbname)
#         except ConnectionError as error:
#             logging.error(
#                 "Connection error while trying to create InfluxDB database: %s. Waiting for retry...", error)
#             time.sleep(self.db_create_delay)
#             self.init_database()
#         except (InfluxDBServerError, InfluxDBClientError) as error:
#             logging.warning(
#                 "Could not create InfluxDB database. Assuming it already exists: %s", error)
# 
#     def flush(self):
#         """
#         Flush values with writer
#         """
#         if not self.buffer:
#             # Don't do anything when buffer empty
#             return
#         try:
#             self.last_flush_time = time.time()
#             self.writer.write(self.buffer)
#             if self.config.statistics:
#                 self.show_statistics()
#         except (InfluxDBServerError, InfluxDBClientError) as influx_error:
#             logging.error("Error while writing to InfluxDB: %s", influx_error)
#         finally:
#             self.buffer = []
# 
#     def show_statistics(self):
#         """
#         Print performance metrics to stdout
#         """
#         delta = time.time() - self.start_time
#         msg_per_sec = self.config.buffer_size / delta
#         print("Flushing output buffer. {0:.2f} messages/s".format(msg_per_sec))
#         # Reset timer
#         self.start_time = time.time()
# 
#     def set_reader(self, reader):
#         self.reader = reader
# 
#     def get_reader(self):
#         return self.reader
# 
#     def set_writer(self, writer):
#         self.writer = writer
# 
#     def get_writer(self):
#         return self.writer
# 
#     def get_buffer(self):
#         return self.buffer
# 
#     def get_config(self):
#         return self.config
# 
# Path: kafka_influxdb/encoder/echo_encoder.py
# class Encoder(object):
#     def encode(msg):
# 
# Path: kafka_influxdb/tests/helpers/timeout.py
# def timeout(seconds=5):
#     def decorator(func):
#         def wrapper(*args, **kwargs):
#             process = Process(None, func, None, args, kwargs)
#             process.start()
#             process.join(seconds)
#             if process.is_alive():
#                 process.terminate()
#         return wraps(func)(wrapper)
#     return decorator


# Path: kafka_influxdb/tests/test_worker.py
import unittest
import random
import time
from mock import Mock
from kafka_influxdb.worker import Worker
from kafka_influxdb.encoder import echo_encoder
from kafka_influxdb.tests.helpers.timeout import timeout
from requests.exceptions import ConnectionError
from mock import patch

        self.buffer_timeout = buffer_timeout
        self.kafka_topic = "test"
        self.influxdb_dbname = "mydb"
        self.statistics = False


class DummyReader(object):
    """
    A reader that yields dummy messages
    """

    def __init__(self, messages, num_messages):
        self.messages = messages
        self.num_messages = num_messages

    def read(self):
        for i in range(self.num_messages):
            yield random.choice(self.messages)
        # Simulate keyboard interrupt by user to stop consuming
        raise KeyboardInterrupt


class TimeoutReader(object):
    """
    A reader that writes half the messages then times out
    """

    def __init__(self, message, num_messages, timeout):
        self.message = message
        self.num_messages = num_messages

        self.timeout = timeout