Fixed tests for new changes

Author: Akash Verma

ingestion/tests/unit/topology/database/test_burstiq_client.py

@@ -31,6 +31,8 @@ def setUp(self):
             username="test_user",
             password="test_password",
             realmName="test_realm",
+            biqSdzName="test_sdz",
+            biqCustomerName="test_customer",
         )
 
         # Mock authentication response
@@ -49,6 +51,8 @@ def test_authentication_success(self, mock_post):
         mock_post.return_value = mock_response
 
         client = BurstIQClient(self.config)
+        # Trigger authentication explicitly
+        client.test_authenticate()
 
         # Verify authentication was called correctly
         self.assertEqual(mock_post.call_count, 1)
@@ -70,8 +74,6 @@ def test_authentication_success(self, mock_post):
         # Verify token is set
         self.assertIsNotNone(client.access_token)
         self.assertIsNotNone(client.token_expires_at)
-        self.assertEqual(client.customer_name, "test_realm")
-        self.assertEqual(client.sdz_name, "test_sdz")
 
     @patch("metadata.ingestion.source.database.burstiq.client.requests.post")
     def test_authentication_failure(self, mock_post):
@@ -80,8 +82,9 @@ def test_authentication_failure(self, mock_post):
         mock_response.raise_for_status.side_effect = Exception("401 Unauthorized")
         mock_post.return_value = mock_response
 
+        client = BurstIQClient(self.config)
         with self.assertRaises(Exception) as context:
-            BurstIQClient(self.config)
+            client.test_authenticate()
 
         self.assertIn("Failed to authenticate with BurstIQ", str(context.exception))
 
@@ -247,6 +250,9 @@ def test_token_refresh(self, mock_request, mock_datetime, mock_post):
 
         client = BurstIQClient(self.config)
 
+        # Trigger initial authentication
+        client.test_authenticate()
+
         # Manually set token to expired
         client.token_expires_at = mock_now - timedelta(seconds=1)
 
@@ -345,6 +351,8 @@ def test_custom_auth_server_url(self, mock_post):
             username="test_user",
             password="test_password",
             realmName="custom_realm",
+            biqSdzName="custom_sdz",
+            biqCustomerName="custom_customer",
         )
 
         mock_response = Mock()
@@ -353,6 +361,8 @@ def test_custom_auth_server_url(self, mock_post):
         mock_post.return_value = mock_response
 
         client = BurstIQClient(custom_config)
+        # Trigger authentication
+        client.test_authenticate()
 
         # Verify custom auth server URL was used
         call_args = mock_post.call_args

ingestion/tests/unit/topology/database/test_burstiq_connection.py

@@ -30,6 +30,8 @@ def setUp(self):
             username="test_user",
             password="test_password",
             realmName="test_realm",
+            biqSdzName="test_sdz",
+            biqCustomerName="test_customer",
         )
 
         self.mock_token_response = {
@@ -49,8 +51,8 @@ def test_get_connection(self, mock_post):
 
         # Verify client was created
         self.assertIsNotNone(client)
-        self.assertEqual(client.username, "test_user")
-        self.assertEqual(client.realm_name, "test_realm")
+        self.assertEqual(client.config.username, "test_user")
+        self.assertEqual(client.config.realmName, "test_realm")
 
     @patch("metadata.ingestion.source.database.burstiq.client.requests.post")
     @patch("metadata.ingestion.source.database.burstiq.client.requests.request")
@@ -197,8 +199,9 @@ def test_connection_invalid_credentials(self, mock_post):
         mock_response.raise_for_status.side_effect = Exception("401 Unauthorized")
         mock_post.return_value = mock_response
 
-        # Should raise exception
+        # Should raise exception when authentication is attempted
+        client = get_connection(self.config)
         with self.assertRaises(Exception) as context:
-            get_connection(self.config)
+            client.test_authenticate()
 
         self.assertIn("Failed to authenticate", str(context.exception))

ingestion/tests/unit/topology/database/test_burstiq_lineage.py

@@ -67,6 +67,8 @@ def test_lineage_source_initialization(self, mock_post):
                     "username": "test_user",
                     "password": "test_password",
                     "realmName": "test_realm",
+                    "biqSdzName": "test_sdz",
+                    "biqCustomerName": "test_customer",
                 }
             },
             "sourceConfig": {"config": {"type": "DatabaseLineage"}},
@@ -118,6 +120,8 @@ def test_get_table_entity_not_found(self, mock_post):
                     "username": "test_user",
                     "password": "test_password",
                     "realmName": "test_realm",
+                    "biqSdzName": "test_sdz",
+                    "biqCustomerName": "test_customer",
                 }
             },
             "sourceConfig": {"config": {"type": "DatabaseLineage"}},
@@ -158,6 +162,8 @@ def test_process_edge_missing_tables(self, mock_post):
                     "username": "test_user",
                     "password": "test_password",
                     "realmName": "test_realm",
+                    "biqSdzName": "test_sdz",
+                    "biqCustomerName": "test_customer",
                 }
             },
             "sourceConfig": {"config": {"type": "DatabaseLineage"}},
@@ -208,6 +214,8 @@ def test_lineage_iteration_success(self, mock_request, mock_post):
                     "username": "test_user",
                     "password": "test_password",
                     "realmName": "test_realm",
+                    "biqSdzName": "test_sdz",
+                    "biqCustomerName": "test_customer",
                 }
             },
             "sourceConfig": {"config": {"type": "DatabaseLineage"}},
@@ -252,6 +260,8 @@ def test_lineage_iteration_error_handling(self, mock_request, mock_post):
                     "username": "test_user",
                     "password": "test_password",
                     "realmName": "test_realm",
+                    "biqSdzName": "test_sdz",
+                    "biqCustomerName": "test_customer",
                 }
             },
             "sourceConfig": {"config": {"type": "DatabaseLineage"}},

ingestion/tests/unit/topology/database/test_burstiq_metadata.py

@@ -35,6 +35,8 @@ def setUp(self):
             username="test_user",
             password="test_password",
             realmName="test_realm",
+            biqSdzName="test_sdz",
+            biqCustomerName="test_customer",
         )
 
         # Sample dictionary data
@@ -272,13 +274,23 @@ def test_table_constraints_unique(self):
 
     def test_table_constraints_foreign_key(self):
         """Test extraction of foreign key constraints from referenced dictionaries"""
+        from unittest.mock import patch
         from metadata.ingestion.source.database.burstiq.metadata import Burstiqsource
 
         source = Mock(spec=Burstiqsource)
         source.get_table_constraints = Burstiqsource.get_table_constraints.__get__(
             source
         )
 
+        # Mock metadata and context for FQN building
+        source.metadata = Mock()
+        source.context = Mock()
+        context_data = Mock()
+        context_data.database_service = "test_service"
+        context_data.database = "test_db"
+        context_data.database_schema = "test_schema"
+        source.context.get.return_value = context_data
+
         # Create dictionary with foreign key reference
         dictionary = BurstIQDictionary(
             name="visit",
@@ -292,14 +304,21 @@ def test_table_constraints_foreign_key(self):
             indexes=[],
         )
 
-        # Get constraints
-        constraints = source.get_table_constraints(dictionary)
-
-        # Verify foreign key constraint
-        self.assertIsNotNone(constraints)
-        self.assertEqual(len(constraints), 1)
-        self.assertEqual(constraints[0].constraintType, ConstraintType.FOREIGN_KEY)
-        self.assertEqual(constraints[0].columns, ["patient_id"])
+        # Mock fqn.build to return a test FQN
+        with patch('metadata.ingestion.source.database.burstiq.metadata.fqn.build') as mock_fqn_build:
+            mock_fqn_build.return_value = "test_service.test_db.test_schema.patient.patient_id"
+
+            # Get constraints
+            constraints = source.get_table_constraints(dictionary)
+
+            # Verify foreign key constraint
+            self.assertIsNotNone(constraints)
+            self.assertEqual(len(constraints), 1)
+            self.assertEqual(constraints[0].constraintType, ConstraintType.FOREIGN_KEY)
+            self.assertEqual(constraints[0].columns, ["patient_id"])
+            # Verify the FQN was built and used (it's wrapped in FullyQualifiedEntityName)
+            self.assertEqual(len(constraints[0].referredColumns), 1)
+            self.assertEqual(constraints[0].referredColumns[0].root, "test_service.test_db.test_schema.patient.patient_id")
 
     def test_table_name_extraction(self):
         """Test table name and type extraction from dictionaries"""