add test for column badge

Signed-off-by: adoami <adoami@chanzuckerberg.com>

tests/unit/proxy/test_neo4j_proxy.py

@@ -412,7 +412,7 @@ def test_delete_owner(self) -> None:
             self.assertEqual(mock_run.call_count, 1)
             self.assertEqual(mock_commit.call_count, 1)
 
-    def test_add_badge(self) -> None:
+    def test_add_table_badge(self) -> None:
         with patch.object(GraphDatabase, 'driver') as mock_driver:
             mock_session = MagicMock()
             mock_driver.return_value.session.return_value = mock_session
@@ -433,6 +433,28 @@ def test_add_badge(self) -> None:
             self.assertEqual(mock_run.call_count, 3)
             self.assertEqual(mock_commit.call_count, 1)
 
+    def test_add_column_badge(self) -> None:
+        with patch.object(GraphDatabase, 'driver') as mock_driver:
+            mock_session = MagicMock()
+            mock_driver.return_value.session.return_value = mock_session
+
+            mock_transaction = MagicMock()
+            mock_session.begin_transaction.return_value = mock_transaction
+
+            mock_run = MagicMock()
+            mock_transaction.run = mock_run
+            mock_commit = MagicMock()
+            mock_transaction.commit = mock_commit
+
+            neo4j_proxy = Neo4jProxy(host='DOES_NOT_MATTER', port=0000)
+            neo4j_proxy.add_badge(id='dummy_uri',
+                                  badge_name='hive',
+                                  column_name='dummy_column',
+                                  resource_type=ResourceType.Column)
+            # we call neo4j twice in add_tag call
+            self.assertEqual(mock_run.call_count, 3)
+            self.assertEqual(mock_commit.call_count, 1)
+
     def test_add_tag(self) -> None:
         with patch.object(GraphDatabase, 'driver') as mock_driver:
             mock_session = MagicMock()