Add more tests for async run client

 * Directly inherit sync client's methods in async mode

Author: md-mq

cli/polyaxon/_client/run.py

@@ -3232,7 +3232,7 @@ def push_offline_run(
 class AsyncRunClient(RunClient):
     _IS_ASYNC = True
 
-    async def _areset_agent(self, agent: Optional[str] = None):
+    async def _reset_agent(self, agent: Optional[str] = None):
         if agent:
             agent = await self.client.agents_v1.get_agent(self.owner, agent)
             self.settings.agent = V1SettingsCatalog(
@@ -3244,7 +3244,7 @@ async def _areset_agent(self, agent: Optional[str] = None):
             self.settings.namespace = agent.namespace
             self.settings.artifacts_store = None
 
-    async def _ause_agent_host(self):
+    async def _use_agent_host(self):
         if self.settings.agent and self.settings.agent.url:
             await self.areset_client(
                 host=self.settings.agent.url,
@@ -3469,7 +3469,7 @@ async def get_events(
     ):
         if not self.settings:
             await self.refresh_data()
-        await self._ause_agent_host()
+        await self._use_agent_host()
 
         params = get_streams_params(self.artifacts_store)
         return await self.client.runs_v1.get_run_events(
@@ -3495,7 +3495,7 @@ async def get_multi_run_events(
     ):
         if not self.settings:
             await self.refresh_data()
-        await self._ause_agent_host()
+        await self._use_agent_host()
 
         params = get_streams_params(self.artifacts_store)
         return await self.client.runs_v1.get_multi_run_events(

cli/tests/test_client/test_async_run_client.py

@@ -528,3 +528,80 @@ async def test_load_offline_run_is_sync_only():
 
     with pytest.raises(PolyaxonClientException):
         await AsyncRunClient.load_offline_run("/tmp/run")
+
+
+@pytest.mark.asyncio
+async def test_restart_routes_to_copy_when_copy_flag_set():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    response = make_run(name="restarted")
+    sdk_client.runs_v1.copy_run = AsyncMock(return_value=response)
+    sdk_client.runs_v1.restart_run = AsyncMock()
+    client = make_client(sdk_client)
+
+    result = await client.restart(copy=True, name="restarted")
+
+    assert result is response
+    assert sdk_client.runs_v1.copy_run.call_count == 1
+    assert sdk_client.runs_v1.restart_run.call_count == 0
+    assert "async_req" not in sdk_client.runs_v1.copy_run.call_args[1]
+
+
+@pytest.mark.asyncio
+async def test_restart_routes_to_restart_when_copy_not_set():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    response = make_run(name="restarted")
+    sdk_client.runs_v1.restart_run = AsyncMock(return_value=response)
+    sdk_client.runs_v1.copy_run = AsyncMock()
+    client = make_client(sdk_client)
+
+    result = await client.restart(name="restarted")
+
+    assert result is response
+    assert sdk_client.runs_v1.restart_run.call_count == 1
+    assert sdk_client.runs_v1.copy_run.call_count == 0
+    assert "async_req" not in sdk_client.runs_v1.restart_run.call_args[1]
+
+
+@pytest.mark.asyncio
+async def test_resume_awaits_resume_run_with_built_body():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    response = make_run(name="resumed")
+    sdk_client.runs_v1.resume_run = AsyncMock(return_value=response)
+    client = make_client(sdk_client)
+
+    result = await client.resume(name="resumed")
+
+    assert result is response
+    assert sdk_client.runs_v1.resume_run.call_count == 1
+    assert "async_req" not in sdk_client.runs_v1.resume_run.call_args[1]
+
+
+@pytest.mark.asyncio
+async def test_log_succeeded_awaits_create_run_status_via_log_status():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    sdk_client.runs_v1.create_run_status = AsyncMock(return_value=None)
+    client = make_client(sdk_client)
+    client._run_data = make_run(status=V1Statuses.RUNNING)
+
+    await client.log_succeeded()
+
+    assert sdk_client.runs_v1.create_run_status.call_count == 1
+    body = sdk_client.runs_v1.create_run_status.call_args[1]["body"]
+    assert body["condition"].type == V1Statuses.SUCCEEDED
+
+
+@pytest.mark.asyncio
+async def test_log_end_status_short_circuits_when_already_done():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    sdk_client.runs_v1.create_run_status = AsyncMock()
+    client = make_client(sdk_client)
+    client._run_data = make_run(status=V1Statuses.SUCCEEDED)
+
+    await client.log_failed()
+
+    assert sdk_client.runs_v1.create_run_status.call_count == 0