test_database.py

'''
Test cases related to the contents of database.py:
- Connection build-up and data storage
'''

import unittest
import pandas as pd
from database import DatabaseHandler

class TestDatabaseHandler(unittest.TestCase):

    def setUp(self):
        # Create a temporary in-memory SQLite database and generate test data
        self.db = DatabaseHandler(":memory:")
        self.test_df = pd.DataFrame({'x': [1.0, 2.0], 'y': [3.0, 4.0]})

    def tearDown(self):
        self.db.close()

    def test_save_data_succeeds(self):
        # Try to save table
        self.db.save_data(self.test_df, 'test_table')

        # Verify that the data was actually saved
        with self.db.engine.connect() as conn:
            result = pd.read_sql('SELECT * FROM test_table', conn)

        self.assertEqual(len(result), 2)
        self.assertIn('x', result.columns)
        self.assertIn('y', result.columns)

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