refactor(cli): make internal agent client explicit

 * issue #1151

Author: md-mq

cli/polyaxon/_client/client.py

@@ -23,6 +23,7 @@
 )
 from polyaxon._sdk.async_client.api_client import AsyncApiClient
 from polyaxon._sdk.sync_client.api_client import ApiClient
+from polyaxon.exceptions import PolyaxonClientException
 
 if TYPE_CHECKING:
     from polyaxon._schemas.client import ClientConfig
@@ -72,6 +73,7 @@ def __init__(
         config: Optional["ClientConfig"] = None,
         token: Optional[str] = None,
         is_async: bool = False,
+        is_internal: bool = False,
     ):
         self._config = config or settings.CLIENT_CONFIG
         token = token or self._config.token
@@ -83,6 +85,7 @@ def __init__(
             self._config.token = token
 
         self.is_async = is_async
+        self.is_internal = is_internal
         self.api_client = self._get_client()
         self._projects_v1 = None
         self._runs_v1 = None
@@ -92,7 +95,6 @@ def __init__(
         self._users_v1 = None
         self._versions_v1 = None
         self._agents_v1 = None
-        self._internal_agents_v1 = None
         self._queues_v1 = None
         self._service_accounts_v1 = None
         self._presets_v1 = None
@@ -104,21 +106,18 @@ def __init__(
         self._organizations_v1 = None
 
     def _get_client(self):
+        if self.is_internal:
+            headers = self.config.get_internal_header()
+            if self.is_async:
+                return AsyncApiClient(self.config.async_internal_sdk_config, **headers)
+            return ApiClient(self.config.internal_sdk_config, **headers)
+
         if self.is_async:
             return AsyncApiClient(
                 self.config.async_sdk_config, **self.config.client_header
             )
         return ApiClient(self.config.sdk_config, **self.config.client_header)
 
-    def _get_internal_client(self):
-        if self.is_async:
-            return AsyncApiClient(
-                self.config.async_internal_sdk_config, **self.config.get_internal_header()
-            )
-        return ApiClient(
-            self.config.internal_sdk_config, **self.config.get_internal_header()
-        )
-
     def reset(self):
         self._projects_v1 = None
         self._runs_v1 = None
@@ -128,7 +127,6 @@ def reset(self):
         self._users_v1 = None
         self._versions_v1 = None
         self._agents_v1 = None
-        self._internal_agents_v1 = None
         self._queues_v1 = None
         self._service_accounts_v1 = None
         self._presets_v1 = None
@@ -180,12 +178,6 @@ def agents_v1(self):
             self._agents_v1 = AgentsV1Api(self.api_client)
         return self._agents_v1
 
-    @property
-    def internal_agents_v1(self):
-        if not self._internal_agents_v1:
-            self._internal_agents_v1 = AgentsV1Api(self._get_internal_client())
-        return self._internal_agents_v1
-
     @property
     def queues_v1(self):
         if not self._queues_v1:

cli/polyaxon/_runner/agent/client.py

@@ -13,12 +13,14 @@ def __init__(
         owner: Optional[str] = None,
         agent_uuid: Optional[str] = None,
         client: Optional[PolyaxonClient] = None,
+        internal_client: Optional[PolyaxonClient] = None,
         is_async: bool = False,
     ):
         self.owner = owner
         self.agent_uuid = agent_uuid
         self.is_async = is_async
         self._client = client
+        self._internal_client = internal_client
 
     @property
     def client(self):
@@ -27,6 +29,16 @@ def client(self):
         self._client = PolyaxonClient(is_async=self.is_async)
         return self._client
 
+    @property
+    def internal_client(self):
+        if self._internal_client:
+            return self._internal_client
+        self._internal_client = PolyaxonClient(
+            is_async=self.is_async,
+            is_internal=True,
+        )
+        return self._internal_client
+
     @property
     def _is_managed(self):
         return self.owner is not None and self.agent_uuid is not None
@@ -68,7 +80,7 @@ def cron_agent(self):
         )
 
     def collect_agent_data(self, namespace: str):
-        return self.client.internal_agents_v1.collect_agent_data(
+        return self.internal_client.agents_v1.collect_agent_data(
             owner=self.owner,
             uuid=self.agent_uuid,
             namespace=namespace,

cli/tests/test_client/test_polyaxon_client.py

@@ -1,6 +1,8 @@
 import pytest
 import tempfile
 
+from mock import patch
+
 from polyaxon import settings
 from polyaxon._client.client import PolyaxonClient
 from polyaxon._constants.globals import NO_AUTH
@@ -10,9 +12,10 @@
     ProjectsV1Api,
     RunsV1Api,
     UsersV1Api,
-    VersionsV1Api,
+    VersionsV1Api, AgentsV1Api,
 )
 from polyaxon._utils.test_utils import BaseTestCase
+from polyaxon.exceptions import PolyaxonClientException
 
 
 @pytest.mark.client_mark
@@ -24,6 +27,7 @@ def setUp(self):
     def test_client_services(self):
         settings.AUTH_CONFIG.token = None
         client = PolyaxonClient(token=None)
+        assert client.is_internal is False
         assert client.config.token is None
 
         assert isinstance(client.config, ClientConfig)
@@ -34,6 +38,36 @@ def test_client_services(self):
         assert isinstance(client.runs_v1, RunsV1Api)
         assert isinstance(client.users_v1, UsersV1Api)
 
+    def test_internal_client_services(self):
+        client = PolyaxonClient(is_internal=True)
+        assert client.is_internal is True
+        assert isinstance(client.agents_v1, AgentsV1Api)
+
+    def test_sync_internal_client_uses_internal_config(self):
+        settings.CLIENT_CONFIG.token = "token"
+        with patch("polyaxon._client.client.ApiClient") as api_client:
+            PolyaxonClient(is_internal=True)
+
+        sdk_config = api_client.call_args.args[0]
+        assert sdk_config.api_key_prefix["ApiKey"] == "InternalToken"
+        assert (
+            api_client.call_args.kwargs
+            == settings.CLIENT_CONFIG.get_internal_header()
+        )
+
+    def test_async_internal_client_uses_async_internal_config(self):
+        settings.CLIENT_CONFIG.token = "token"
+        with patch("polyaxon._client.client.AsyncApiClient") as api_client:
+            PolyaxonClient(is_async=True, is_internal=True)
+
+        sdk_config = api_client.call_args.args[0]
+        assert sdk_config.api_key_prefix["ApiKey"] == "InternalToken"
+        assert sdk_config.connection_pool_maxsize == 100
+        assert (
+            api_client.call_args.kwargs
+            == settings.CLIENT_CONFIG.get_internal_header()
+        )
+
     def test_from_config(self):
         settings.CLIENT_CONFIG.host = "localhost"
         client = PolyaxonClient(config=ClientConfig())

cli/tests/test_runner/test_base_agent.py

@@ -127,3 +127,30 @@ def test_init_agent_component(self, register):
         assert isinstance(agent.executor, MagicMock)
         assert isinstance(agent.client, AgentClient)
         assert register.call_count == 1
+
+    def test_agent_client_uses_internal_client_for_collect_agent_data(self):
+        public_client = MagicMock()
+        internal_client = MagicMock()
+        client = AgentClient(
+            owner="foo",
+            agent_uuid="uuid",
+            client=public_client,
+            internal_client=internal_client,
+        )
+
+        client.collect_agent_data(namespace="default")
+
+        assert public_client.agents_v1.collect_agent_data.call_count == 0
+        internal_client.agents_v1.collect_agent_data.assert_called_once_with(
+            owner="foo",
+            uuid="uuid",
+            namespace="default",
+        )
+
+    @patch("polyaxon._runner.agent.client.PolyaxonClient")
+    def test_agent_client_creates_internal_client_with_internal_mode(self, client_cls):
+        client = AgentClient(owner="foo", agent_uuid="uuid", is_async=True)
+
+        _ = client.internal_client
+
+        client_cls.assert_called_once_with(is_async=True, is_internal=True)