cdplib/cdplib/unit_tests/TestMongodbHandler.py

import unittest
import sys
import os
from pymongo import MongoClient
sys.path.append(os.getcwd())
from cdplib.log import Log
from cdplib.db_handlers.MongodbHandler import MongodbHandler

class TestMongodbHandler(unittest.TestCase):

    def setUp(self):
        database_url = "mongodb://{0}:{1}@{2}:{3}"\
                           .format('root',
                                   'oebb',
                                   'localhost',
                                   27017)
        self.database_name = 'test_database'
        self.first_collection_name = 'first_test_collection'
        self.second_collection_name = 'second_test_collection'
        self.mongodb_handler = MongodbHandler(database_name=self.database_name, database_url=database_url)
        self.client = MongoClient(database_url) 
        self.database = self.client[self.database_name]
        self.valid_schema_path = os.path.join('.', 'cdplib', 'unit_tests', "valid_test_schema.json")
        self.invalid_schema_path = os.path.join('.', 'cdplib', 'unit_tests', "invalid_test_schema.json")
        self.valid_input = {
                        "test_value_string": "test_value",
                        "test_value_double": 2.4,
                        "test_value_double_array": [1.4, 1.6, 3.5]
                        }
        self.invalid_input = {
                        "test_value_string": 1,
                        "test_value_double": "Wrong value",
                        "test_value_double_array": [1.4, 1.6, 3.5]
                        }


    def test_A_set_database(self):
        '''
        Checks that the database doesnt exist.
        Creates the database.
        Create a collection in the database.
        Check that the database now exists.
        Check that an assertionerror is thrown then an integer value is sent to the method.
        '''
        self.assertFalse(self.database_name in self.client.list_database_names())
        #self.test_B_create_collection()
        #self.assertTrue(self.database_name in self.client.list_database_names())
        self.assertRaises(AssertionError, lambda:self.mongodb_handler.set_database(123))

    def test_B_create_collection(self):
        '''
        Checks that the collection doesnt exist
        Creates the collection
        Checks that the collection now exists
        Checks that an Exception is thrown when an integervalue is given to the method.
        '''
        self.assertFalse(self.first_collection_name in self.database.list_collection_names())
        self.mongodb_handler.create_collection(self.first_collection_name)
        self.mongodb_handler.create_collection(self.second_collection_name)
        self.assertTrue(self.first_collection_name in self.database.list_collection_names())
        self.assertRaises(Exception, lambda:self.mongodb_handler.create_collection(123))

    def test_C_set_collection_schema(self):
        '''
        Sets a schema for the collection
        Tries to set a schema which has syntax errors
        '''
        self.mongodb_handler.set_collection_schema(self.first_collection_name, self.valid_schema_path)
        self.assertRaises(Exception, lambda:set_collection_schema(self.first_collection_name, self.invalid_schema_path))

    def test_D_insert_data_into_collection(self):
        '''
        Inserts data into the collection
        Tries to insert data which doesnt comform to the schema.
        '''
        self.mongodb_handler.insert_data_into_collection(self.valid_input, self.first_collection_name)
        self.assertRaises(Exception, lambda:set_collection_schemaself.mongodb_handler.insert_data_into_collection(self.invalid_input, self.first_collection_name))
        
    def test_E_query_data_and_generate_dataframe(self):
        '''
        Fetch data and confirms thats it is the same as was entered into the database
        Do the same with more specific query
        '''
        self.assertEqual(self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name).to_dict()['test_value_double'][0], self.valid_input['test_value_double'])
        self.assertEqual(self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name, 'test_value_string', 'test_value').to_dict()['test_value_double'][0], self.valid_input['test_value_double'])
    
    def test_F_aggregate_data_and_generate_dataframe(self):
        '''
        Make an aggregation call
        Make sure its the same data as was entered into the database
        '''
        aggregation_pipeline = [
                                { '$match': {}}
                                ]
        self.assertEqual(self.mongodb_handler.aggregate_data_and_generate_dataframe(self.first_collection_name, aggregation_pipeline).to_dict()['test_value_double'][0], self.valid_input['test_value_double'])
    
    def test_G_update_data_in_collection(self):
        '''
        Fetch data from database
        Make sure its the same as was entered into the database
        Update a value in the data
        Fetch the new data
        Make sure that it is different from the original data.
        '''
        original_value = self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name).to_dict()['test_value_string'][0]
        self.assertEqual(original_value, self.valid_input['test_value_string'])
        self.mongodb_handler.update_data_in_collection('test_value_string', 'new_test_value', self.first_collection_name, 'test_value_string', 'test_value', create_if_not_exist=False)
        new_value =  self.mongodb_handler.query_data_and_generate_dataframe(self.first_collection_name).to_dict()['test_value_string'][0]
        self.assertNotEqual(original_value, new_value)

    def test_H_create_index(self):
        '''
        Create a index in the collection
        Make sure that the index exists in the collection indexes
        '''
        index = 'test_value_string'
        self.mongodb_handler.create_index(self.first_collection_name, index)
        self.assertTrue(index in list(self.database[self.first_collection_name].index_information().keys()))
    
    def test_Y_drop_collection(self):
        '''
        Checks that the collection exists
        Drops the collection
        Checks that the collection doesn't exist anymore
        '''
        self.assertTrue(self.first_collection_name in self.database.list_collection_names())
        self.mongodb_handler.drop_collection(self.first_collection_name)
        self.assertFalse(self.first_collection_name in self.database.list_collection_names())
    
    def test_Z_drop_database(self):
        '''
        Checks that the database exists
        Drops the database
        Checks that the database doesn't exist anymore
        '''
        self.assertTrue(self.database_name in self.client.list_database_names())
        self.mongodb_handler.drop_database()
        self.assertFalse(self.database_name in self.client.list_database_names())

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