Repository: microsoft/MLOS

Issue:
Generate the HTTP header fields for authorization directly in AzureAuthService
That is, move `AzureVMService._get_headers()` functionality to `AzureAuthService` and make it generic enough to expose via `SupportsAuth` interface.


Discussion:

Resolving patch:
diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py b/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py
--- a/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py
+++ b/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py
@@ -50,7 +50,10 @@ def __init__(self,
         """
         super().__init__(
             config, global_config, parent,
-            self.merge_methods(methods, [self.get_access_token])
+            self.merge_methods(methods, [
+                self.get_access_token,
+                self.get_auth_headers,
+            ])
         )
 
         # This parameter can come from command line as strings, so conversion is needed.
@@ -119,3 +122,9 @@ def get_access_token(self) -> str:
             self._access_token = res.token
             _LOG.info("Got new accessToken. Expiration time: %s", self._token_expiration_ts)
         return self._access_token
+
+    def get_auth_headers(self) -> dict:
+        """
+        Get the authorization part of HTTP headers for REST API calls.
+        """
+        return {"Authorization": "Bearer " + self.get_access_token()}
diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_services.py b/mlos_bench/mlos_bench/services/remote/azure/azure_services.py
--- a/mlos_bench/mlos_bench/services/remote/azure/azure_services.py
+++ b/mlos_bench/mlos_bench/services/remote/azure/azure_services.py
@@ -207,7 +207,7 @@ def _get_headers(self) -> dict:
         """
         assert self._parent is not None and isinstance(self._parent, SupportsAuth), \
             "Authorization service not provided. Include service-auth.jsonc?"
-        return {"Authorization": "Bearer " + self._parent.get_access_token()}
+        return self._parent.get_auth_headers()
 
     @staticmethod
     def _extract_arm_parameters(json_data: dict) -> dict:
diff --git a/mlos_bench/mlos_bench/services/types/authenticator_type.py b/mlos_bench/mlos_bench/services/types/authenticator_type.py
--- a/mlos_bench/mlos_bench/services/types/authenticator_type.py
+++ b/mlos_bench/mlos_bench/services/types/authenticator_type.py
@@ -15,8 +15,6 @@ class SupportsAuth(Protocol):
     Protocol interface for authentication for the cloud services.
     """
 
-    # pylint: disable=too-few-public-methods
-
     def get_access_token(self) -> str:
         """
         Get the access token for cloud services.
@@ -26,3 +24,13 @@ def get_access_token(self) -> str:
         access_token : str
             Access token.
         """
+
+    def get_auth_headers(self) -> dict:
+        """
+        Get the authorization part of HTTP headers for REST API calls.
+
+        Returns
+        -------
+        access_header : dict
+            HTTP header containing the access token.
+        """


Test patch:
diff --git a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_auth_service.py b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_auth_service.py
--- a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_auth_service.py
+++ b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_auth_service.py
@@ -26,8 +26,14 @@ def __init__(self, config: Optional[Dict[str, Any]] = None,
                  methods: Union[Dict[str, Callable], List[Callable], None] = None):
         super().__init__(
             config, global_config, parent,
-            self.merge_methods(methods, [self.get_access_token])
+            self.merge_methods(methods, [
+                self.get_access_token,
+                self.get_auth_headers,
+            ])
         )
 
     def get_access_token(self) -> str:
         return "TOKEN"
+
+    def get_auth_headers(self) -> dict:
+        return {"Authorization": "Bearer " + self.get_access_token()}