Spaces:
Runtime error
Runtime error
| from unittest import TestCase, main | |
| from db_operation import db_operator | |
| from datetime import datetime | |
| class test_db_operation(TestCase): | |
| def setUp(self) -> None: | |
| self.db_operator = db_operator('sqlite:///local_db.db') | |
| self.db_operator.delete_all_stocks() | |
| self.stock1 = { | |
| "ticker": 'AAPL', | |
| "weight": 1.0, | |
| "display_name": 'Apple Inc.', | |
| "date": datetime(2021, 1, 1) | |
| } | |
| self.stock2 = { | |
| "ticker": 'MSFT', | |
| "weight": 1.0, | |
| "display_name": 'Microsoft Corporation', | |
| "date": datetime(2021, 1, 1) | |
| } | |
| def test_insert(self): | |
| stock1 = { | |
| "ticker": 'AAPL', | |
| "weight": 1.0, | |
| "display_name": 'Apple Inc.', | |
| "date": datetime(2021, 1, 1) | |
| } | |
| self.db_operator.add_stock(stock1) | |
| retrieved_stock = self.db_operator.get_stocks_between(datetime(2021, 1, 1), datetime(2021, 1, 1))[0] | |
| self.assertEqual(retrieved_stock.ticker, 'AAPL') | |
| self.assertEqual(retrieved_stock.weight, 1.0) | |
| self.assertEqual(retrieved_stock.display_name, 'Apple Inc.') | |
| self.assertEqual(retrieved_stock.date, datetime(2021, 1, 1)) | |
| def test_delete(self): | |
| self.db_operator.add_stock(self.stock1) | |
| self.db_operator.add_stock(self.stock2) | |
| self.db_operator.delete_stocks_between( | |
| datetime(2021, 1, 1), | |
| datetime(2021, 1, 1)) | |
| retrieved_stocks = self.db_operator.get_stocks_between( | |
| datetime(2021, 1, 1), | |
| datetime(2021, 1, 1)) | |
| self.assertEqual(len(retrieved_stocks), 0) | |
| def test_query_window_1d(self): | |
| # insert 2 stocks between 2021-01-01 and 2021-01-01 every hour | |
| for i in range(24): | |
| self.stock1['date'] = datetime(2021, 1, 1, i) | |
| self.stock2['date'] = datetime(2021, 1, 1, i) | |
| self.db_operator.add_stock(self.stock1) | |
| self.db_operator.add_stock(self.stock2) | |
| # insert two on 2021-01-02 | |
| self.stock1['date'] = datetime(2021, 1, 2) | |
| self.stock2['date'] = datetime(2021, 1, 2) | |
| self.db_operator.add_stock(self.stock1) | |
| self.db_operator.add_stock(self.stock2) | |
| # query 1d | |
| retrieved_stocks = self.db_operator.get_stocks_between( | |
| datetime(2021, 1, 1), | |
| datetime(2021, 1, 2)) | |
| self.assertEqual(len(retrieved_stocks), 50) | |
| def test_query_window_12h(self): | |
| # insert 2 stocks every hour between 2021-01-01 and 2021-01-01 | |
| for i in range(24): | |
| self.stock1['date'] = datetime(2021, 1, 1, i) | |
| self.stock2['date'] = datetime(2021, 1, 1, i) | |
| self.db_operator.add_stock(self.stock1) | |
| self.db_operator.add_stock(self.stock2) | |
| # query 12h | |
| retrieved_stocks = self.db_operator.get_stocks_between( | |
| datetime(2021, 1, 1, 0), | |
| datetime(2021, 1, 1, 12)) | |
| self.assertEqual(len(retrieved_stocks), 26) | |
| # self.assertTrue(False) | |
| def test_query_window_1h(self): | |
| # insert 2 stocks every mins between 2021-01-01 and 2021-01-01 | |
| for i in range(60): | |
| self.stock1['date'] = datetime(2021, 1, 1, 0, i) | |
| self.stock2['date'] = datetime(2021, 1, 1, 0, i) | |
| self.db_operator.add_stock(self.stock1) | |
| self.db_operator.add_stock(self.stock2) | |
| # query 1h | |
| retrieved_stocks = self.db_operator.get_stocks_between( | |
| datetime(2021, 1, 1, 0), | |
| datetime(2021, 1, 1, 1)) | |
| self.assertEqual(len(retrieved_stocks), 120) | |
| def test_query_window_30m(self): | |
| # insert 2 stocks every 1 between 2021-01-01-00:00 and 2021-01-01-00:20 | |
| for i in range(20): | |
| self.stock1['date'] = datetime(2021, 1, 1, 0, i) | |
| self.stock2['date'] = datetime(2021, 1, 1, 0, i) | |
| self.db_operator.add_stock(self.stock1) | |
| self.db_operator.add_stock(self.stock2) | |
| # query 30m | |
| retrieved_stocks = self.db_operator.get_stocks_between( | |
| datetime(2021, 1, 1, 0), | |
| datetime(2021, 1, 1, 0, 30)) | |
| self.assertEqual(len(retrieved_stocks), 40) | |
| if __name__ == '__main__': | |
| main() | |