Write unit test for mongodb handler

cdplib/db_handlers/MongodbHandler.py

@@ -84,19 +84,37 @@ class MongodbHandler:
 
         self._database_name = database_name
 
-    def set_database(self, database_name):
+    def set_database(self, database_name: str):
+        '''
+        :param str database_name: Name of the database.
+        '''
+        assert(isinstance(database_name, str)),\
+            "Parameter 'database_name' must be a string type"
+
+        if database_name not in self._client.list_database_names():
+            self._log.info(('Database: {} didnt exist, it will be created for you once a collection is created in it').format(database_name))
         self._database = self._client[database_name]
 
+
     def drop_database(self):
         '''
         '''
-        self._client.drop_database(self._database_name)
+        try:
+            self._client.drop_database(self._database_name)
+        except Exception as error:
+            self._log.log_and_raise_error(('Couldnt drop the database. Error: {}').format(error))
+    
 
     def drop_collection(self, collection_name: str):
         '''
         '''
-        self._database[collection_name].drop()
-
+        try:
+            return_message = self._database.drop_collection(collection_name)
+            if 'errmsg' in return_message:
+                raise Exception(return_message['errmsg'])
+        except Exception as error:
+            self._log.log_and_raise_error(('Couldnt drop the collection {}. Error: {}').format(collection_name, error))
+    
     def create_index(self, collection_name: str, key: str, direction: (str, int)='text'):
         '''
         :param str collection_name: name on the collection for which the schema will be set.
@@ -108,12 +126,15 @@ class MongodbHandler:
 
         assert(isinstance(collection_name, str)),\
             "Parameter 'collection_name' must be a string type"
-        assert(isinstance(key, (str, int)) & key in allowed_directions),\
-            "Parameter 'key' must be a string or integer type and must be one of these values: 1, -1, '2d', 'geoHaystack', '2dsphere', 'hashed', 'text' "
+        assert(isinstance(key, (str, int))),\
+            "Parameter 'key' must be a string or integer type"
+        assert(direction in allowed_directions),\
+            "Parameter 'key' must be one of these values: 1, -1, '2d', 'geoHaystack', '2dsphere', 'hashed', 'text' "
         assert(isinstance(direction, str)),\
             "Parameter 'direction' must be a string type"
             
-        self._database[collection_name].create_index([key, direction])
+        self._database[collection_name].create_index([(key, direction)], name=key)
+        #collection.create_index([('field_i_want_to_index', pymongo.TEXT)], name='search_index', default_language='english')
 
     def set_collection_schema(self, collection_name: str, schema_path: str,
                               validation_level: str = 'moderate',validation_action: str = 'error'):
@@ -135,7 +156,6 @@ class MongodbHandler:
         parse_obj = ParseJsonSchema(schema_paths=schema_path)
 
         schema = parse_obj.read_schema_and_parse_for_mongodb(schema_path)
-
         command = {
                     'collMod': collection_name,
                     'validator': {
@@ -162,13 +182,13 @@ class MongodbHandler:
         if collection_name not in self._database.list_collection_names():
             try:
                 self._log.info(("Collection '{}' has been created").format(collection_name))
-                return self._database.create_collection(collection_name)
+                self._database.create_collection(collection_name)
 
             except Exception as error:
                 self._log.log_and_raise_error(('An error occured while creating the new collection: {}. \nError: {}').format(collection_name, error))
         else:
             self._log.info(("Collection '{}' already exists").format(collection_name))
-            return self._database[collection_name]
+            self._database[collection_name]
 
     def insert_data_into_collection(self, data: (dict, list, np.ndarray, pd.DataFrame, pd.Series),
                                     collection_name: str,
@@ -226,7 +246,7 @@ class MongodbHandler:
         self.set_collection_schema(collection_name=collection_name, schema_path=schema_path)
 
     def query_data_and_generate_dataframe(self, collection_name: str, attribute: str = None,
-                                          attribute_value: str = None, comparison_operator: str = '$eq'):
+                                          attribute_value: str = None, comparison_operator: str = '$eq', index = None):
         '''
 
         '''
@@ -239,25 +259,25 @@ class MongodbHandler:
         except Exception as error:
             self._log.log_and_raise_error(('An error occured trying to query data from {}, with query {}: {}:{}. \nError:{}').format(collection_name, attribute_value, comparison_operator, attribute_value, error))
 
-        return self.convert_mongo_data_into_dataframe(data)
+        return self.convert_mongo_data_into_dataframe(data, index)
 
     def aggregate_data_and_generate_dataframe(self, collection_name: str, aggregation_pipeline: list):
 
         try:
             data = self._database[collection_name].aggregate(pipeline=aggregation_pipeline, allowDiskUse=True)
-
         except Exception as error:
             self._log.log_and_raise_error(('A problem occured when aggregating the collection {} with the pipeline {}. \nError: {}').format(collection_name, aggregation_pipeline, error))
 
         return self.convert_mongo_data_into_dataframe(data)
 
-    def convert_mongo_data_into_dataframe(self, data) -> pd.DataFrame():
+    def convert_mongo_data_into_dataframe(self, data, index: str = None) -> pd.DataFrame():
 
         data = list(data)
         try:
             if len(data)> 0:
                 df = pd.DataFrame(data)
-                df.set_index('radsatznummer', inplace=True)
+                if index is not None:
+                    df.set_index(index, inplace=True)
                 return df
             else:
                 self._log.warning(('No data for the query was found').format())

cdplib/unit_tests/TestMongodbHandler.py

@@ -0,0 +1,142 @@
+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', 'test_value', 'test_value_string', 'new_test_value', self.first_collection_name)
+        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())
+        self.assertRaises(Exception, lambda:self.mongodb_handler.drop_collection(self.first_collection_name))
+    
+    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()

cdplib/unit_tests/invalid_test_schema.json

@@ -0,0 +1,19 @@
+{
+    "bsonType": "object",
+    "properties": {
+
+        "test_value_string": {
+            "bsonType": "string",
+        },
+        "test_value_double": {
+            "bsonType": "double",
+        },
+        "test_value_double_array": {
+            "bsonType": "array",
+            "items": {
+                "bsonType": "double"
+            }
+        }
+    }
+}
+}

cdplib/unit_tests/valid_test_schema.json

@@ -0,0 +1,17 @@
+{
+    "bsonType": "object",
+    "properties": {
+        "test_value_string": {
+            "bsonType": "string"
+        },
+        "test_value_double": {
+            "bsonType": "double"
+        },
+        "test_value_double_array": {
+            "bsonType": "array",
+            "items": {
+                "bsonType": "double"
+            }
+        }
+    }
+}