cleanup mock logic

metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py

@@ -84,13 +84,14 @@ def __init__(
         tenant_id: str,
         metadata_api_timeout: int,
     ):
-        self.__access_token: Optional[str] = None
-        self.__access_token_expiry_time: Optional[datetime] = None
-        self.__tenant_id = tenant_id
+        self._access_token: Optional[str] = None
+        self._access_token_expiry_time: Optional[datetime] = None
+
+        self._tenant_id = tenant_id
         # Test connection by generating access token
         logger.info(f"Trying to connect to {self._get_authority_url()}")
         # Power-Bi Auth (Service Principal Auth)
-        self.__msal_client = msal.ConfidentialClientApplication(
+        self._msal_client = msal.ConfidentialClientApplication(
             client_id,
             client_credential=client_secret,
             authority=DataResolverBase.AUTHORITY + tenant_id,
@@ -168,18 +169,18 @@ def _get_app(
         pass
 
     def _get_authority_url(self):
-        return f"{DataResolverBase.AUTHORITY}{self.__tenant_id}"
+        return f"{DataResolverBase.AUTHORITY}{self._tenant_id}"
 
     def get_authorization_header(self):
         return {Constant.Authorization: self.get_access_token()}
 
-    def get_access_token(self):
-        if self.__access_token is not None and not self._is_access_token_expired():
-            return self.__access_token
+    def get_access_token(self) -> str:
+        if self._access_token is not None and not self._is_access_token_expired():
+            return self._access_token
 
         logger.info("Generating PowerBi access token")
 
-        auth_response = self.__msal_client.acquire_token_for_client(
+        auth_response = self._msal_client.acquire_token_for_client(
             scopes=[DataResolverBase.SCOPE]
         )
 
@@ -193,24 +194,24 @@ def get_access_token(self):
 
         logger.info("Generated PowerBi access token")
 
-        self.__access_token = "Bearer {}".format(
+        self._access_token = "Bearer {}".format(
             auth_response.get(Constant.ACCESS_TOKEN)
         )
         safety_gap = 300
-        self.__access_token_expiry_time = datetime.now() + timedelta(
+        self._access_token_expiry_time = datetime.now() + timedelta(
             seconds=(
                 max(auth_response.get(Constant.ACCESS_TOKEN_EXPIRY, 0) - safety_gap, 0)
             )
         )
 
-        logger.debug(f"{Constant.PBIAccessToken}={self.__access_token}")
+        logger.debug(f"{Constant.PBIAccessToken}={self._access_token}")
 
-        return self.__access_token
+        return self._access_token
 
     def _is_access_token_expired(self) -> bool:
-        if not self.__access_token_expiry_time:
+        if not self._access_token_expiry_time:
             return True
-        return self.__access_token_expiry_time < datetime.now()
+        return self._access_token_expiry_time < datetime.now()
 
     def get_dashboards(self, workspace: Workspace) -> List[Dashboard]:
         """

metadata-ingestion/tests/integration/powerbi/test_powerbi.py

@@ -31,19 +31,18 @@
 
 def mock_msal_cca(*args, **kwargs):
     class MsalClient:
-        call_num = 0
-        token: Dict[str, Any] = {
-            "access_token": "dummy",
-        }
+        def __init__(self):
+            self.call_num = 0
+            self.token: Dict[str, Any] = {
+                "access_token": "dummy",
+            }
 
-        @staticmethod
-        def acquire_token_for_client(*args, **kwargs):
-            MsalClient.call_num += 1
-            return MsalClient.token
+        def acquire_token_for_client(self, *args, **kwargs):
+            self.call_num += 1
+            return self.token
 
-        @staticmethod
-        def reset():
-            MsalClient.call_num = 0
+        def reset(self):
+            self.call_num = 0
 
     return MsalClient()
 
@@ -739,81 +738,84 @@ def test_workspace_container(
     )
 
 
-@mock.patch("msal.ConfidentialClientApplication", side_effect=mock_msal_cca)
 def test_access_token_expiry_with_long_expiry(
-    mock_msal: MagicMock,
     pytestconfig: pytest.Config,
     tmp_path: str,
     mock_time: datetime.datetime,
     requests_mock: Any,
 ) -> None:
     register_mock_api(pytestconfig=pytestconfig, request_mock=requests_mock)
 
-    pipeline = Pipeline.create(
-        {
-            "run_id": "powerbi-test",
-            "source": {
-                "type": "powerbi",
-                "config": {
-                    **default_source_config(),
+    mock_msal = mock_msal_cca()
+
+    with mock.patch("msal.ConfidentialClientApplication", return_value=mock_msal):
+        pipeline = Pipeline.create(
+            {
+                "run_id": "powerbi-test",
+                "source": {
+                    "type": "powerbi",
+                    "config": {
+                        **default_source_config(),
+                    },
                 },
-            },
-            "sink": {
-                "type": "file",
-                "config": {
-                    "filename": f"{tmp_path}/powerbi_access_token_mces.json",
+                "sink": {
+                    "type": "file",
+                    "config": {
+                        "filename": f"{tmp_path}/powerbi_access_token_mces.json",
+                    },
                 },
-            },
-        }
-    )
+            }
+        )
 
     # for long expiry, the token should only be requested once.
-    MsalClient.token = {
+    mock_msal.token = {
         "access_token": "dummy2",
         "expires_in": 3600,
     }
+    mock_msal.reset()
 
-    MsalClient.reset()
     pipeline.run()
     # We expect the token to be requested twice (once for AdminApiResolver and one for RegularApiResolver)
-    assert MsalClient.call_num == 2
+    assert mock_msal.call_num == 2
 
 
-@mock.patch("msal.ConfidentialClientApplication", side_effect=mock_msal_cca)
 def test_access_token_expiry_with_short_expiry(
-    mock_msal: MagicMock,
     pytestconfig: pytest.Config,
     tmp_path: str,
     mock_time: datetime.datetime,
     requests_mock: Any,
 ) -> None:
     register_mock_api(pytestconfig=pytestconfig, request_mock=requests_mock)
 
-    pipeline = Pipeline.create(
-        {
-            "run_id": "powerbi-test",
-            "source": {
-                "type": "powerbi",
-                "config": {
-                    **default_source_config(),
+    mock_msal = mock_msal_cca()
+    with mock.patch("msal.ConfidentialClientApplication", return_value=mock_msal):
+        pipeline = Pipeline.create(
+            {
+                "run_id": "powerbi-test",
+                "source": {
+                    "type": "powerbi",
+                    "config": {
+                        **default_source_config(),
+                    },
                 },
-            },
-            "sink": {
-                "type": "file",
-                "config": {
-                    "filename": f"{tmp_path}/powerbi_access_token_mces.json",
+                "sink": {
+                    "type": "file",
+                    "config": {
+                        "filename": f"{tmp_path}/powerbi_access_token_mces.json",
+                    },
                 },
-            },
-        }
-    )
+            }
+        )
 
     # for short expiry, the token should be requested when expires.
-    MsalClient.token = {
+    mock_msal.token = {
         "access_token": "dummy",
         "expires_in": 0,
     }
+    mock_msal.reset()
+
     pipeline.run()
-    assert MsalClient.call_num > 2
+    assert mock_msal.call_num > 2
 
 
 def dataset_type_mapping_set_to_all_platform(pipeline: Pipeline) -> None: