feat: add compute engine module with acquire token from gcp

serpens/compute_engine.py

@@ -0,0 +1,25 @@
+import logging
+import requests
+
+logger = logging.getLogger(__name__)
+
+METADATA_HEADERS = {"Metadata-Flavor": "Google"}
+METADATA_VM_IDENTITY_URL = (
+    "http://metadata.google.internal/computeMetadata/v1/"
+    "instance/service-accounts/default/identity?"
+    "audience={audience}&format={format}&licenses={licenses}"
+)
+
+
+def acquire_token(audience, format="standard", licenses=True):
+    url = METADATA_VM_IDENTITY_URL.format(audience=audience, format=format, licenses=licenses)
+
+    response = requests.get(url, headers=METADATA_HEADERS)
+
+    logger.debug(f"Response status: {response.status_code}")
+    logger.debug(f"Response content: {response.content}")
+
+    if not response.ok:
+        return None
+
+    return response.text

tests/test_compute_engine.py

@@ -0,0 +1,39 @@
+from http import HTTPStatus
+import unittest
+from unittest.mock import patch
+
+import compute_engine
+from requests import Response
+
+
+class TestComputeEngineGCP(unittest.TestCase):
+    def setUp(self):
+        self.patch_requests = patch("compute_engine.requests")
+        self.mock_requests = self.patch_requests.start()
+
+        self.audience = "http://www.example.com"
+
+    def tearDown(self) -> None:
+        self.patch_requests.stop()
+
+    def test_acquire_token(self):
+        token = "eyJhbGciOiJSUzI1Ni.eyJhdWQiOiJodHRwczov.b25hd3MuY29tIi"
+
+        mock_response = Response()
+        mock_response.status_code = HTTPStatus.OK.value
+        mock_response._content = token.encode()
+
+        self.mock_requests.get.return_value = mock_response
+
+        acquired_token = compute_engine.acquire_token(audience=self.audience)
+
+        self.assertEqual(acquired_token, token)
+
+    def test_acquire_token_fail(self):
+        mock_response = Response()
+        mock_response.status_code = HTTPStatus.INTERNAL_SERVER_ERROR.value
+
+        self.mock_requests.get.return_value = mock_response
+
+        response = compute_engine.acquire_token(audience=self.audience)
+        self.assertIsNone(response)