implement insert and delete functions for node_provider_lookup table

aviate-labs · Nov 14, 2023 · b0d3310 · b0d3310
1 parent ca42814
commit b0d3310
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 23 deletions.
diff --git a/node_monitor/node_provider_db.py b/node_monitor/node_provider_db.py
@@ -163,7 +163,7 @@ def _execute1(self, sql: str, params: Tuple[Any, ...]) -> List[Tuple[Any, ...]]:
         return result
 
 
-    def execute_write(self, sql: str, params: Tuple[Any, ...]) -> None:
+    def _execute_write(self, sql: str, params: Tuple[Any, ...]) -> None:
         conn = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(sql, params)
@@ -241,6 +241,27 @@ def get_node_providers_as_dict(self) -> Dict[Principal, str]:
         rows = self._execute("SELECT * FROM node_provider_lookup", ())
         lookupd = {row['node_provider_id']: row['node_provider_name'] for row in rows}
         return lookupd
+
+
+    def insert_node_provider(self, node_provider: ic_api.NodeProvider) -> None:
+        """Inserts a NodeProvider object into node_provider_lookup"""
+        query = """
+            INSERT INTO node_provider_lookup (
+                node_provider_id,
+                node_provider_name
+            ) VALUES (%s, %s)
+        """
+        params = (node_provider.principal_id, node_provider.display_name)
+        self._execute_write(query, params)
+
+
+    def delete_node_provider(self, node_provider_id: Principal) -> None:
+        """Deletes a record in node_proivder_lookup by node_provider_id"""
+        query = """
+            DELETE FROM node_provider_lookup
+            WHERE node_provider_id = %s
+        """
+        self._execute_write(query, (node_provider_id,))
 
 
     def close(self) -> None:

diff --git a/tests/test_node_provider_db.py b/tests/test_node_provider_db.py
@@ -2,6 +2,7 @@
 from devtools import debug
 
 from node_monitor.node_provider_db import NodeProviderDB
+from node_monitor.ic_api import NodeProvider
 from node_monitor.node_monitor_helpers.sql_constants import \
     DELETE_SUBSCRIBER, INSERT_SUBSCRIBER
 from tests.conftest import cached
@@ -73,36 +74,21 @@ def test_validate_schema_node_provider_lookup():
 
 
 @pytest.mark.db
-def test_insert_subscriber_crud():
+def test_insert_and_delete_node_provider():
 
     # Insert new node providers
     new_node_providers = cached["new_node_providers"]  
     for node_provider in new_node_providers.node_providers:
-        params = (node_provider.principal_id, False, False, 
-                    False, node_provider.display_name, False)
-        node_provider_db.execute_write(INSERT_SUBSCRIBER, params)
+        node_provider_db.insert_node_provider(node_provider)
 
     # Check new node providers were added properly
-    subs = node_provider_db.get_subscribers_as_dict()
-    assert subs['test-dummy-principal-1'] == \
-        {'node_provider_id': 'test-dummy-principal-1', 'notify_on_status_change': False, 'notify_email': False,
-         'notify_slack': False, 'node_provider_name': 'Node Provider A', 'notify_telegram': False,}
-    assert subs['test-dummy-principal-2'] == \
-        {'node_provider_id': 'test-dummy-principal-2', 'notify_on_status_change': False, 'notify_email': False,
-         'notify_slack': False, 'node_provider_name': 'Node Provider B', 'notify_telegram': False,}
-
-    # Check overwriting existing records occurs properly 
-    params = ('test-dummy-principal-1', False, False, 
-              False, 'Node Provider C', False)
-    node_provider_db.execute_write(INSERT_SUBSCRIBER, params)
-    subs = node_provider_db.get_subscribers_as_dict()
-    assert subs['test-dummy-principal-1'] == \
-        {'node_provider_id': 'test-dummy-principal-1', 'notify_on_status_change': False, 'notify_email': False,
-         'notify_slack': False, 'node_provider_name': 'Node Provider C', 'notify_telegram': False,}
+    node_providers = node_provider_db.get_node_providers_as_dict()
+    assert node_providers['test-dummy-principal-1'] == "Node Provider A"
+    assert node_providers['test-dummy-principal-2'] == "Node Provider B"
 
     # Delete subscribers
-    node_provider_db.execute_write(DELETE_SUBSCRIBER, ('test-dummy-principal-1',))
-    node_provider_db.execute_write(DELETE_SUBSCRIBER, ('test-dummy-principal-2',))
+    node_provider_db.delete_node_provider('test-dummy-principal-1')
+    node_provider_db.delete_node_provider('test-dummy-principal-2')
     subs = node_provider_db.get_subscribers_as_dict()
     assert 'test-dummy-principal-1' not in subs
     assert 'test-dummy-principal-2' not in subs