add base class for mongo dao

Author: ash-raph

dao/_base_mongo.py

@@ -0,0 +1,26 @@
+from pymongo import MongoClient
+
+from sample_code.settings import (
+    ARC_MONGO_AUTHMECHANISM,
+    ARC_MONGO_AUTHSOURCE,
+    ARC_MONGO_READ_PREFERENCE,
+)
+
+
+class BaseMongoDAO:
+    def __init__(
+        self,
+        mongoServers: str,
+        mongoReplicaset: str,
+        username: str,
+        password: str,
+        database: str,
+    ) -> None:
+        mongo_uri = f"mongodb://{username}:{password}@{mongoServers}"
+        self.client = MongoClient(
+            mongo_uri,
+            replicaSet=mongoReplicaset,
+            authSource=ARC_MONGO_AUTHSOURCE,
+            readPreference=ARC_MONGO_READ_PREFERENCE,
+            authMechanism=ARC_MONGO_AUTHMECHANISM,
+        )[database]

dao/audit.py

@@ -1,32 +1,27 @@
 from datetime import datetime
 
-from pymongo import MongoClient
-
+from sample_code.dao._base_mongo import BaseMongoDAO
 from sample_code.settings import (
-    ARC_MONGO_AUTHMECHANISM,
-    ARC_MONGO_AUTHSOURCE,
-    ARC_MONGO_READ_PREFERENCE,
     AUDIT_COLLECTION,
+    AUDIT_DATABASE,
+    AUDIT_PASSWORD,
+    AUDIT_USERNAME,
 )
 
 
-class AuditDAO:
+class AuditDAO(BaseMongoDAO):
     def __init__(
         self,
         mongoServers: str,
         mongoReplicaset: str,
-        username: str,
-        password: str,
-        database: str,
     ) -> None:
-        mongo_uri = f"mongodb://{username}:{password}@{mongoServers}"
-        self.client = MongoClient(
-            mongo_uri,
-            replicaSet=mongoReplicaset,
-            authSource=ARC_MONGO_AUTHSOURCE,
-            readPreference=ARC_MONGO_READ_PREFERENCE,
-            authMechanism=ARC_MONGO_AUTHMECHANISM,
-        )[database]
+        super().__init__(
+            username=AUDIT_USERNAME,
+            password=AUDIT_PASSWORD,
+            database=AUDIT_DATABASE,
+            mongoServers=mongoServers,
+            mongoReplicaset=mongoReplicaset,
+        )
 
     def run_aggregation_query(collection, query, **kwargs):
         return collection.aggregate(query, **kwargs)

dao/usage.py

@@ -1,33 +1,25 @@
 from datetime import datetime
 
-from pymongo import DESCENDING, MongoClient
+from pymongo import DESCENDING
 from pymongo.collection import Collection
 
-from sample_code.settings import (
-    ARC_MONGO_AUTHMECHANISM,
-    ARC_MONGO_AUTHSOURCE,
-    ARC_MONGO_READ_PREFERENCE,
-    COLLECTION,
-)
+from sample_code.dao._base_mongo import BaseMongoDAO
+from sample_code.settings import COLLECTION, DATABASE, PASSWORD, USERNAME
 
 
-class UsageDAO:
+class UsageDAO(BaseMongoDAO):
     def __init__(
         self,
         mongoServers: str,
         mongoReplicaset: str,
-        username: str,
-        password: str,
-        database: str,
     ) -> None:
-        mongo_uri = f"mongodb://{username}:{password}@{mongoServers}"
-        self.client = MongoClient(
-            mongo_uri,
-            replicaSet=mongoReplicaset,
-            authSource=ARC_MONGO_AUTHSOURCE,
-            readPreference=ARC_MONGO_READ_PREFERENCE,
-            authMechanism=ARC_MONGO_AUTHMECHANISM,
-        )[database]
+        super().__init__(
+            username=USERNAME,
+            password=PASSWORD,
+            database=DATABASE,
+            mongoServers=mongoServers,
+            mongoReplicaset=mongoReplicaset,
+        )
 
     def run_query(
         collection: Collection,

main.py

@@ -7,21 +7,15 @@
 from sample_code.dao.reporting import ReportDAO
 from sample_code.dao.usage import UsageDAO
 from sample_code.settings import (
-    AUDIT_DATABASE,
-    AUDIT_PASSWORD,
     AUDIT_REPLICASET,
     AUDIT_SERVER,
-    AUDIT_USERNAME,
-    DATABASE,
-    PASSWORD,
     REPLICASET_A,
     REPLICASET_B,
     REPLICASET_C,
     REPORTING_AULDATALEAK_TABLENAME,
     SERVER_A,
     SERVER_B,
     SERVER_C,
-    USERNAME,
 )
 
 
@@ -32,33 +26,21 @@ def __init__(self) -> None:
         self.auditClient = AuditDAO(
             mongoServers=AUDIT_SERVER,
             mongoReplicaset=AUDIT_REPLICASET,
-            username=AUDIT_USERNAME,
-            password=AUDIT_PASSWORD,
-            database=AUDIT_DATABASE,
         )
 
         self.usageClient_A = UsageDAO(
             mongoServers=SERVER_A,
             mongoReplicaset=REPLICASET_A,
-            username=USERNAME,
-            password=PASSWORD,
-            database=DATABASE,
         )
 
         self.usageClient_B = UsageDAO(
             mongoServers=SERVER_B,
             mongoReplicaset=REPLICASET_B,
-            username=USERNAME,
-            password=PASSWORD,
-            database=DATABASE,
         )
 
         self.usageClient_C = UsageDAO(
             mongoServers=SERVER_C,
             mongoReplicaset=REPLICASET_C,
-            username=USERNAME,
-            password=PASSWORD,
-            database=DATABASE,
         )
 
     def get_auldata_subscribers(self, auditRangeStart, auditRangeEnd):

tests/conftest.py

@@ -4,8 +4,8 @@
 
 
 @pytest.fixture(autouse=True)
-def mock_audit_mongo_client():
-    mock_mongo_client = patch("sample_code.dao.audit.MongoClient")
+def mock_mongo_client():
+    mock_mongo_client = patch("sample_code.dao._base_mongo.MongoClient")
     yield mock_mongo_client.start()
     mock_mongo_client.stop()
 
@@ -17,13 +17,6 @@ def mock_usage_mongo_run_agg_query():
     mock_run_query.stop()
 
 
-@pytest.fixture(autouse=True)
-def mock_usage_mongo_client():
-    mock_mongo_client = patch("sample_code.dao.usage.MongoClient")
-    yield mock_mongo_client.start()
-    mock_mongo_client.stop()
-
-
 @pytest.fixture(autouse=True)
 def mock_usage_mongo_run_query():
     mock_run_query = patch("sample_code.dao.usage.UsageDAO.run_query")

tests/test_audit.py

@@ -6,9 +6,7 @@
 def test_get_subscribers(mock_usage_mongo_run_agg_query):
     startDate = date.today() - timedelta(days=1)
     endDate = date.today()
-    auditClient = AuditDAO(
-        "mongo-server.com", "mongo-replicaset", "username", "secret_password", "db-name"
-    )
+    auditClient = AuditDAO("mongo-server.com", "mongo-replicaset")
 
     auditClient.get_subscribers(startDate, endDate)
     assert mock_usage_mongo_run_agg_query.called_once()

tests/test_usage.py

@@ -6,9 +6,7 @@
 def test_get_subscriber_usage(mock_usage_mongo_run_query):
     effDate = date.today() - timedelta(days=1)
     expDate = date.today()
-    usageClient = UsageDAO(
-        "mongo-server.com", "mongo-replicaset", "username", "secret_password", "db-name"
-    )
+    usageClient = UsageDAO("mongo-server.com", "mongo-replicaset")
 
     usageClient.get_subscriber_usage("1", effDate, expDate)
     assert mock_usage_mongo_run_query.called_once()