import unittest import pandas as pd import random from typing import List from analytic_types.data_bucket import DataBucket from tests.test_dataset import create_list_of_timestamps class TestBucket(unittest.TestCase): def test_receive_data(self): bucket = DataBucket() data_val = list(range(6)) timestamp_list = create_list_of_timestamps(len(data_val)) for val in data_val: bucket.receive_data(get_pd_dataframe([val], [1523889000000 + val])) for idx, row in bucket.data.iterrows(): self.assertEqual(data_val[idx], row['value']) self.assertEqual(timestamp_list[idx], row['timestamp']) def test_drop_data(self): bucket = DataBucket() data_val = list(range(10)) timestamp_list = create_list_of_timestamps(len(data_val)) bucket.receive_data(get_pd_dataframe(data_val, timestamp_list)) bucket.drop_data(5) expected_data = data_val[5:] expected_timestamp = timestamp_list[5:] self.assertEqual(expected_data, bucket.data['value'].tolist()) self.assertEqual(expected_timestamp, bucket.data['timestamp'].tolist()) if __name__ == '__main__': unittest.main() def get_pd_dataframe(value: List[int], timestamp: List[int]) -> pd.DataFrame: if len(value) != len(timestamp): raise ValueError(f'len(value) should be equal to len(timestamp)') return pd.DataFrame({ 'value': value, 'timestamp': timestamp })