code: import unittest
  import pandas as pd
  import numpy as np
  import context
  
  from fastbt.datasource import DataSource
  import talib
  
  class TestDataSource(unittest.TestCase):
  
      def setUp(self):        
          df = pd.read_csv('tests/data/sample.csv', parse_dates=['timestamp'])
          self.ds = DataSource(data=df)
  
      def test_data(self):
          self.assertEqual(self.ds.data.iloc[20,1], 'five')
          self.assertEqual(self.ds.data.iloc[14,3], 112)
          self.assertEqual(self.ds.data.iloc[24,7], 10.54)
  
      def test_data_without_sort(self):
          df = pd.read_csv('tests/data/sample.csv', parse_dates=['timestamp'])
          self.ds = DataSource(data=df, sort=False)
          self.assertEqual(self.ds.data.iloc[9,4], 999)
          self.assertEqual(self.ds.data.iloc[24,6], 41688)
          self.assertEqual(self.ds.data.at[4, 'close'], 10.6)
  
      def test_initialize_case(self):
          df = pd.read_csv('tests/data/sample.csv', parse_dates=['timestamp'])
          df.columns = [x.upper() for x in df.columns]
          self.assertEqual(df.columns[0], 'TIMESTAMP')
          self.ds = DataSource(data=df)
          self.assertEqual(self.ds.data.columns[0], 'timestamp')
  
      def test_initialize_column_rename(self):
          df = pd.read_csv('tests/data/sample.csv', parse_dates=['timestamp'])
          df.columns = ['TS', 'TRADINGSYMBOL', 'OPEN', 'HIGH', 'LOW',
                      'CLOSE', 'VOLUME', 'PREVCLOSE']
          self.ds = DataSource(data=df, timestamp='TS', symbol='TRADINGSYMBOL')
          self.assertEqual(self.ds.data.columns[0], 'timestamp')
          self.assertEqual(self.ds.data.columns[1], 'symbol')
  
      def test_add_lag(self):
          length = len(self.ds.data)
          idx = pd.IndexSlice
          self.ds.add_lag(on='close')
          self.ds.add_lag(on='volume', period=2)
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          self.assertEqual(d.at[idx['2018-01-04', 'one'], 'lag_close_1'], 11)
          self.assertEqual(d.at[idx['2018-01-06', 'six'], 'lag_volume_2'], 86014)
          self.assertEqual(len(self.ds.data.columns), 10)
          self.assertEqual(len(self.ds.data), length)
  
      def test_add_lag_column_rename(self):
          idx = pd.IndexSlice
          self.ds.add_lag(on='close')
          self.ds.add_lag(on='close', col_name='some_col')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          self.assertEqual(d.at[idx['2018-01-04', 'one'], 'lag_close_1'], 11)
          self.assertEqual(d.at[idx['2018-01-04', 'one'], 'some_col'], 11)
          self.assertEqual(d.at[idx['2018-01-05', 'three'], 'some_col'], 109)
  
      def test_add_pct_change(self):
          idx = pd.IndexSlice
          self.ds.add_pct_change(on='close')
          self.ds.add_pct_change(on='close', period=2)
          self.ds.add_pct_change(on='close', period=2, col_name='new_col')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(R(d.at[idx['2018-01-05', 'three'], 'chg_close_1']), -0.07)
          self.assertEqual(R(d.at[idx['2018-01-06', 'five'], 'chg_close_1']), 0.17)
          self.assertEqual(R(d.at[idx['2018-01-05', 'four'], 'chg_close_2']), 0.05)
          self.assertEqual(R(d.at[idx['2018-01-05', 'four'], 'new_col']), 0.05)
          self.assertEqual(R(d.at[idx['2018-01-03', 'six'], 'new_col']), -0.1)
          self.assertEqual(pd.isna(d.at[idx['2018-01-02', 'one'], 'new_col']), True)
          self.assertEqual(len(self.ds.data.columns), 11)
  
      def test_add_pct_change_lag(self):
          idx = pd.IndexSlice
          self.ds.add_pct_change(on='close', period=2, lag=1)
          self.ds.add_pct_change(on='close', period=1, lag=2)
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(R(d.at[idx['2018-01-04', 'four'], 'chg_close_2']), 0.09)
          self.assertEqual(R(d.at[idx['2018-01-04', 'four'], 'chg_close_1']), 0.01)
          self.assertEqual(R(d.at[idx['2018-01-06', 'three'], 'chg_close_1']), -0.01)
  
  
      def test_add_pct_change_lag_col_name(self):
          idx = pd.IndexSlice
          self.ds.add_pct_change(on='high', period=2, lag=1)
          self.ds.add_pct_change(on='close', period=1, lag=2, col_name='lagged_2')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(R(d.at[idx['2018-01-05', 'six'], 'chg_high_2']), -0.04)
          self.assertEqual(R(d.at[idx['2018-01-04', 'four'], 'lagged_2']), 0.01)
  
      def test_formula_add_col_name(self):
          idx = pd.IndexSlice
          self.ds.add_formula('open+close', 'new_col')
          self.ds.add_formula('volume/close', 'new_col_2')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(R(d.at[idx['2018-01-04', 'four'], 'new_col']), 336)
          self.assertEqual(R(d.at[idx['2018-01-06', 'one'], 'new_col_2']), 77755.77)
  
      def test_formula_case_insensitive(self):
          idx = pd.IndexSlice
          self.ds.add_formula('OPEN+CLOSE', 'new_col')
          self.ds.add_formula('volume/close', 'NEW_COL_2')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(R(d.at[idx['2018-01-04', 'four'], 'new_col']), 336)
          self.assertEqual(R(d.at[idx['2018-01-06', 'one'], 'new_col_2']), 77755.77)
  
  
      def test_formula_calculated_column(self):
          idx = pd.IndexSlice
          self.ds.add_formula('(open+close)*100', 'new_col_1')
          self.ds.add_formula('volume/100', 'new_col_2')
          self.ds.add_formula('new_col_1+new_col_2', 'new_col_3')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(R(d.at[idx['2018-01-06', 'one'], 'new_col_3']), 10190.6)
          self.assertEqual(R(d.at[idx['2018-01-05', 'two'], 'new_col_3']), 200389.97)
  
      def test_rolling_simple(self):
          from pandas import isna
          q = 'symbol == "one"'
          df = pd.read_csv('tests/data/sample.csv', parse_dates=['timestamp']).query(q)
          df['r2'] = df['close'].rolling(2).mean()
          self.ds.add_rolling(2, col_name='r2')
          df2 = self.ds.data.query(q)
          print('RESULT' , df['r2'], df2['r2'])
          for a,b in zip(df['r2'], df2['r2']):
              if not(isna(a)):
                  assert a==b     
  
      def test_rolling_values(self):
          idx = pd.IndexSlice
          self.ds.add_rolling(4, on='volume', function='max')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          R = lambda x: round(x,2)
          self.assertEqual(d.at[idx['2018-01-05', 'five'], 'rol_max_volume_4'], 971704)
          self.assertEqual(d.at[idx['2018-01-05', 'six'], 'rol_max_volume_4'], 195539)
          self.assertEqual(d.at[idx['2018-01-04', 'three'], 'rol_max_volume_4'], 433733)
          # Adding lag and testing
          self.ds.add_rolling(4, on='volume', function='max', lag=1)
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          self.assertEqual(d.at[idx['2018-01-06', 'five'], 'rol_max_volume_4'], 971704)
          self.assertEqual(d.at[idx['2018-01-06', 'six'], 'rol_max_volume_4'], 195539)
          self.assertEqual(d.at[idx['2018-01-05', 'three'], 'rol_max_volume_4'], 433733)
          # Testing for 2 lags and column name
          self.ds.add_rolling(4, on='volume', function='max', lag=2, col_name='check')
          d = self.ds.data.set_index(['timestamp', 'symbol'])
          self.assertEqual(d.at[idx['2018-01-06', 'three'], 'check'], 433733) 
  
      def test_batch(self):
          length = len(self.ds.data)
          batch = [
              {'P': {'on': 'close', 'period': 1, 'lag': 1}},
              {'L': {'on': 'volume', 'period': 1}},
              {'F': {'formula': '(open+close)/2', 'col_name': 'AvgPrice'}},
              {'I': {'indicator': 'SMA', 'period': 3, 'lag': 1, 'col_name': 'SMA3'}},
              {'F': {'formula': 'avgprice + sma3', 'col_name': 'final'}},
              {'R': {'window': 3, 'function': 'mean'}}
          ]
          d = self.ds.batch_process(batch).set_index(['timestamp', 'symbol'])
          self.assertEqual(len(d.columns), 12)
          self.assertEqual(len(self.ds.data.columns), 14)
          self.assertEqual(len(self.ds.data), length)
  
      def test_raise_error_if_not_dataframe(self):
          pass
  
  
  def test_rolling_zscore():
      np.random.seed(100)
      df = pd.DataFrame(np.random.randn(100,4), 
                    columns=['open', 'high', 'low', 'close'])
      df['symbol'] = list('ABCD') * 25
      dates = list(pd.date_range(end='2018-04-25', periods=25)) * 4
      df['timestamp'] = dates
      from fastbt.datasource import DataSource
      ds = DataSource(df)
      ds.add_rolling(on='close', window=5, function='zscore')
      assert ds.data.query('symbol=="A"').iloc[8]['rol_zscore_close_5'].round(2) == 0.12
      assert ds.data.query('symbol=="B"').iloc[-7]['rol_zscore_close_5'].round(2) == 0.17
      assert ds.data.query('symbol=="C"').iloc[-6]['rol_zscore_close_5'].round(2) == -0.48
  
  class TestDataSourceReindex(unittest.TestCase):
  
      def setUp(self): 
          df = pd.DataFrame(np.arange(24).reshape(6, 4),
              columns=['open', 'high', 'low', 'close'])
          df['symbol'] = list('ABCABA')
          df['timestamp'] = [1, 1, 1, 2, 3, 3]
          self.df = df
  
      def test_reindex(self):
          ds = DataSource(self.df)
          ds.reindex([1,2,3])
          assert len(ds.data) == 9
          # Check values
          assert ds.data.set_index(['symbol', 'timestamp']).at[('A', 1), 'open'] == 0
          assert ds.data.set_index(['symbol', 'timestamp']).at[('B', 2), 'close'] == 7
          assert ds.data.set_index(['symbol', 'timestamp']).at[('C', 3), 'high'] == 9
          ds.reindex([1,2,3,4])
          assert len(ds.data) == 12 
  
      def test_reindex_different_fills(self):
          ds = DataSource(self.df)
          ds.reindex([1,2,3], method=None)
          print(ds.data)
          assert pd.isnull(ds.data.set_index(['symbol', 'timestamp']).at[('C', 3), 'high'])
          ds = DataSource(self.df)
          ds.reindex([1,2,3,4], method='bfill') 
          assert ds.data.set_index(['symbol', 'timestamp']).at[('B', 2), 'close'] == 19
  
  class TestDataSourceTALIB(unittest.TestCase):
  
      """
      Test TALIB indicators
      """
  
      def setUp(self):       
          self.df = pd.read_csv('tests/data/sample.csv', parse_dates=['timestamp'])
  
      def test_single_symbol(self):
          df = self.df.query('symbol=="one"')
          ds = DataSource(df)
          ds.add_indicator('SMA', period=3, col_name='sma')
          assert len(ds.data) == 6
  
          sma = talib.SMA(df.close.values, timeperiod=3)
          # If both are equal, there should be no differences
          assert (ds.data.sma - sma).sum() == 0
  
  
  
  
apis: pandas.date_range pandas.read_csv numpy.random.seed numpy.random.randn numpy.arange pandas.isna