class HiveConnection(object):

    def __init__(self, host, user, password):
        """Instantiate a HiveConnector object."""
        self.host = host
        self.user = user
        self.password = password
        self.port = 10000
        self.auth_mechanism = 'PLAIN'
        self._connection = self._connect()

    def _connect(self):
        """Start the connection to database."""
        try:
            return connect(host=self.host, port=self.port,
                           user=self.user, password=self.password,
                           auth_mechanism=self.auth_mechanism)
        except TTransportException as error:
            print('Failed attempt to connect')
            self._connect()

    def _disconnect(self):
        """Close connection to database."""
        self._connection.close()

    def hadoop_connection_handler(function):
        """Start a database connection if not already open."""
        @wraps(function)
        def wrapper(inst, *args, **kwargs):
            if not inst._connection:
                inst._connect()
            return function(inst, *args, **kwargs)
        return wrapper

    @hadoop_connection_handler
    def read(self, query):
        """Execute a query to pull the data.

        Args:
            query: [str] Query to pull the data.

        Returns:
            A list of namedtuple (`Row`).
        """
        columns = self._columns(query)
        cursor = self._connection.cursor()
        cursor.execute(query)
        Record = namedtuple("Record", columns)
        data = map(Record._make, cursor.fetchall())

        cursor.close()
        return data

from unittest.mock import patch, MagicMock
from nose.tools import assert_equal, raises

from services.db_connections import HiveConnection


class TestHiveConnection:
    """Integration test suite for HiveConnection class."""

    def setUp(self):
        self.hive = HiveConnection(user='username', password='password', host='myhost.net')

    def test_reconnect(self):
        """If the connection drops, the object should be able to establish a
        new connection.
        """
        query = 'SELECT * FROM database.table 1'
        self.hive._connect = MagicMock()
        self.hive._disconnect()
        self.hive.read(query)
        assert_equal(self.hive._connect.called, True)