from DataCache import DataCache
from exchangeinterface import ExchangeInterface
import unittest
import pandas as pd
import datetime as dt


class TestDataCache(unittest.TestCase):
    def setUp(self):
        # Setup the database connection here
        self.exchanges = ExchangeInterface()
        self.exchanges.connect_exchange(exchange_name='binance', user_name='test_guy', api_keys=None)
        # This object maintains all the cached data. Pass it connection to the exchanges.
        self.data = DataCache(self.exchanges)

        asset, timeframe, exchange = 'BTC/USD', '2h', 'binance'
        self.key1 = f'{asset}_{timeframe}_{exchange}'

        asset, timeframe, exchange = 'ETH/USD', '2h', 'binance'
        self.key2 = f'{asset}_{timeframe}_{exchange}'

    def test_set_cache(self):
        # Tests
        print('Testing set_cache flag not set:')
        self.data.set_cache(data='data', key=self.key1)
        attr = self.data.__getattribute__('cached_data')
        self.assertEqual(attr[self.key1], 'data')

        self.data.set_cache(data='more_data', key=self.key1)
        attr = self.data.__getattribute__('cached_data')
        self.assertEqual(attr[self.key1], 'more_data')

        print('Testing set_cache no-overwrite flag set:')
        self.data.set_cache(data='even_more_data', key=self.key1, do_not_overwrite=True)
        attr = self.data.__getattribute__('cached_data')
        self.assertEqual(attr[self.key1], 'more_data')

    def test_cache_exists(self):
        # Tests
        print('Testing cache_exists() method:')
        self.assertFalse(self.data.cache_exists(key=self.key2))
        self.data.set_cache(data='data', key=self.key1)
        self.assertTrue(self.data.cache_exists(key=self.key1))

    def test_update_candle_cache(self):
        print('Testing update_candle_cache() method:')
        # Initial data
        df_initial = pd.DataFrame({
            'open_time': [1, 2, 3],
            'open': [100, 101, 102],
            'high': [110, 111, 112],
            'low': [90, 91, 92],
            'close': [105, 106, 107],
            'volume': [1000, 1001, 1002]
        })

        # Data to be added
        df_new = pd.DataFrame({
            'open_time': [3, 4, 5],
            'open': [102, 103, 104],
            'high': [112, 113, 114],
            'low': [92, 93, 94],
            'close': [107, 108, 109],
            'volume': [1002, 1003, 1004]
        })

        self.data.set_cache(data=df_initial, key=self.key1)
        self.data.update_candle_cache(more_records=df_new, key=self.key1)

        result = self.data.get_cache(key=self.key1)
        expected = pd.DataFrame({
            'open_time': [1, 2, 3, 4, 5],
            'open': [100, 101, 102, 103, 104],
            'high': [110, 111, 112, 113, 114],
            'low': [90, 91, 92, 93, 94],
            'close': [105, 106, 107, 108, 109],
            'volume': [1000, 1001, 1002, 1003, 1004]
        })

        pd.testing.assert_frame_equal(result, expected)

    def test_update_cached_dict(self):
        print('Testing update_cached_dict() method:')
        self.data.set_cache(data={}, key=self.key1)
        self.data.update_cached_dict(cache_key=self.key1, dict_key='sub_key', data='value')

        cache = self.data.get_cache(key=self.key1)
        self.assertEqual(cache['sub_key'], 'value')

    def test_get_cache(self):
        print('Testing get_cache() method:')
        self.data.set_cache(data='data', key=self.key1)
        result = self.data.get_cache(key=self.key1)
        self.assertEqual(result, 'data')

    def test_get_records_since(self):
        print('Testing get_records_since() method:')
        df_initial = pd.DataFrame({
            'open_time': [1, 2, 3],
            'open': [100, 101, 102],
            'high': [110, 111, 112],
            'low': [90, 91, 92],
            'close': [105, 106, 107],
            'volume': [1000, 1001, 1002]
        })

        self.data.set_cache(data=df_initial, key=self.key1)
        start_datetime = dt.datetime.utcfromtimestamp(2)
        result = self.data.get_records_since(key=self.key1, start_datetime=start_datetime, record_length=60, ex_details=[]).sort_values(by='open_time').reset_index(drop=True)

        expected = pd.DataFrame({
            'open_time': [2, 3],
            'open': [101, 102],
            'high': [111, 112],
            'low': [91, 92],
            'close': [106, 107],
            'volume': [1001, 1002]
        })

        pd.testing.assert_frame_equal(result, expected)


if __name__ == '__main__':
    unittest.main()