trainml 0.5.16__py3-none-any.whl → 1.0.0__py3-none-any.whl

tests/unit/cloudbender/test_nodes_unit.py

@@ -8,6 +8,7 @@ from aiohttp import WSMessage, WSMsgType
 import trainml.cloudbender.nodes as specimen
 from trainml.exceptions import (
     ApiError,
+    NodeError,
     SpecificationError,
     TrainMLException,
 )
@@ -200,3 +201,114 @@ class nodeTests:
             None,
             dict(command="report"),
         )
+
+    @mark.asyncio
+    async def test_node_wait_for_already_at_status(self, node):
+        """Test wait_for returns immediately if already at target status."""
+        node._status = "active"
+        result = await node.wait_for("active")
+        assert result is None
+
+    @mark.asyncio
+    async def test_node_wait_for_invalid_status(self, node):
+        """Test wait_for raises error for invalid status."""
+        with raises(SpecificationError) as exc_info:
+            await node.wait_for("invalid_status")
+        assert "Invalid wait_for status" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_node_wait_for_timeout_validation(self, node):
+        """Test wait_for validates timeout (line 172)."""
+        node._status = "new"  # Set to different status so timeout check runs
+        with raises(SpecificationError) as exc_info:
+            await node.wait_for("active", timeout=25 * 60 * 60)
+        assert "timeout must be less than" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_node_wait_for_success(self, node, mock_trainml):
+        """Test wait_for succeeds when status matches."""
+        node._status = "new"
+        api_response_new = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            rig_uuid="x",
+            status="new",
+        )
+        api_response_active = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            rig_uuid="x",
+            status="active",
+        )
+        mock_trainml._query = AsyncMock(
+            side_effect=[api_response_new, api_response_active]
+        )
+        with patch("trainml.cloudbender.nodes.asyncio.sleep", new_callable=AsyncMock):
+            result = await node.wait_for("active", timeout=10)
+        assert result == node
+        assert node.status == "active"
+
+    @mark.asyncio
+    async def test_node_wait_for_archived_404(self, node, mock_trainml):
+        """Test wait_for handles 404 for archived status."""
+        node._status = "active"
+        api_error = ApiError(404, {"errorMessage": "Not found"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.nodes.asyncio.sleep", new_callable=AsyncMock):
+            await node.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_node_wait_for_error_status(self, node, mock_trainml):
+        """Test wait_for raises error for errored/failed status."""
+        node._status = "new"
+        api_response_errored = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            rig_uuid="x",
+            status="errored",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_errored)
+        with patch("trainml.cloudbender.nodes.asyncio.sleep", new_callable=AsyncMock):
+            with raises(NodeError):
+                await node.wait_for("active", timeout=10)
+
+    @mark.asyncio
+    async def test_node_wait_for_timeout(self, node, mock_trainml):
+        """Test wait_for raises timeout exception."""
+        node._status = "new"
+        api_response_new = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            rig_uuid="x",
+            status="new",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_new)
+        with patch("trainml.cloudbender.nodes.asyncio.sleep", new_callable=AsyncMock):
+            with raises(TrainMLException) as exc_info:
+                await node.wait_for("active", timeout=0.1)
+        assert "Timeout waiting for" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_node_wait_for_api_error_non_404(self, node, mock_trainml):
+        """Test wait_for raises ApiError when not 404 for archived (line 189)."""
+        node._status = "active"
+        api_error = ApiError(500, {"errorMessage": "Server Error"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.nodes.asyncio.sleep", new_callable=AsyncMock):
+            with raises(ApiError):
+                await node.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_node_wait_for_failed_status(self, node, mock_trainml):
+        """Test wait_for raises error for failed status (line 191)."""
+        node._status = "new"
+        api_response_failed = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            rig_uuid="x",
+            status="failed",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_failed)
+        with patch("trainml.cloudbender.nodes.asyncio.sleep", new_callable=AsyncMock):
+            with raises(NodeError):
+                await node.wait_for("active", timeout=10)

tests/unit/cloudbender/test_providers_unit.py

@@ -139,3 +139,99 @@ class providerTests:
         mock_trainml._query.assert_called_once_with(f"/provider/1", "GET")
         assert provider.id == "provider-id-1"
         assert response.id == "provider-id-1"
+
+    def test_provider_status_property(self, provider):
+        """Test provider status property."""
+        provider._status = "ready"
+        assert provider.status == "ready"
+
+    @mark.asyncio
+    async def test_provider_wait_for_already_at_status(self, provider):
+        """Test wait_for returns immediately if already at target status."""
+        provider._status = "ready"
+        result = await provider.wait_for("ready")
+        assert result is None
+
+    @mark.asyncio
+    async def test_provider_wait_for_invalid_status(self, provider):
+        """Test wait_for raises error for invalid status."""
+        with raises(SpecificationError) as exc_info:
+            await provider.wait_for("invalid_status")
+        assert "Invalid wait_for status" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_provider_wait_for_timeout_validation(self, provider):
+        """Test wait_for validates timeout."""
+        with raises(SpecificationError) as exc_info:
+            await provider.wait_for("ready", timeout=25 * 60 * 60)
+        assert "timeout must be less than" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_provider_wait_for_success(self, provider, mock_trainml):
+        """Test wait_for succeeds when status matches."""
+        provider._status = "new"
+        api_response_new = dict(
+            customer_uuid="a",
+            provider_uuid="1",
+            status="new",
+        )
+        api_response_ready = dict(
+            customer_uuid="a",
+            provider_uuid="1",
+            status="ready",
+        )
+        mock_trainml._query = AsyncMock(
+            side_effect=[api_response_new, api_response_ready]
+        )
+        with patch("trainml.cloudbender.providers.asyncio.sleep", new_callable=AsyncMock):
+            result = await provider.wait_for("ready", timeout=10)
+        assert result == provider
+        assert provider.status == "ready"
+
+    @mark.asyncio
+    async def test_provider_wait_for_archived_404(self, provider, mock_trainml):
+        """Test wait_for handles 404 for archived status."""
+        provider._status = "ready"
+        api_error = ApiError(404, {"errorMessage": "Not found"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.providers.asyncio.sleep", new_callable=AsyncMock):
+            await provider.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_provider_wait_for_error_status(self, provider, mock_trainml):
+        """Test wait_for raises error for errored/failed status."""
+        provider._status = "new"
+        api_response_errored = dict(
+            customer_uuid="a",
+            provider_uuid="1",
+            status="errored",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_errored)
+        with patch("trainml.cloudbender.providers.asyncio.sleep", new_callable=AsyncMock):
+            with raises(specimen.ProviderError):
+                await provider.wait_for("ready", timeout=10)
+
+    @mark.asyncio
+    async def test_provider_wait_for_timeout(self, provider, mock_trainml):
+        """Test wait_for raises timeout exception."""
+        provider._status = "new"
+        api_response_new = dict(
+            customer_uuid="a",
+            provider_uuid="1",
+            status="new",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_new)
+        with patch("trainml.cloudbender.providers.asyncio.sleep", new_callable=AsyncMock):
+            with raises(TrainMLException) as exc_info:
+                await provider.wait_for("ready", timeout=0.1)
+        assert "Timeout waiting for" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_provider_wait_for_api_error_non_404(self, provider, mock_trainml):
+        """Test wait_for raises ApiError when not 404 for archived (line 115)."""
+        provider._status = "ready"
+        api_error = ApiError(500, {"errorMessage": "Server Error"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.providers.asyncio.sleep", new_callable=AsyncMock):
+            with raises(ApiError):
+                await provider.wait_for("archived", timeout=10)

tests/unit/cloudbender/test_regions_unit.py

@@ -195,3 +195,109 @@ class regionTests:
         mock_trainml._query.assert_called_once_with(
             "/provider/1/region/a/checkpoint", "POST", None, expected_payload
         )
+
+    @mark.asyncio
+    async def test_region_wait_for_already_at_status(self, region):
+        """Test wait_for returns immediately if already at target status."""
+        region._status = "healthy"
+        result = await region.wait_for("healthy")
+        assert result is None
+
+    @mark.asyncio
+    async def test_region_wait_for_invalid_status(self, region):
+        """Test wait_for raises error for invalid status."""
+        with raises(SpecificationError) as exc_info:
+            await region.wait_for("invalid_status")
+        assert "Invalid wait_for status" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_region_wait_for_timeout_validation(self, region):
+        """Test wait_for validates timeout (line 135)."""
+        region._status = "new"  # Set to different status so timeout check runs
+        with raises(SpecificationError) as exc_info:
+            await region.wait_for("healthy", timeout=25 * 60 * 60)
+        assert "timeout must be less than" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_region_wait_for_success(self, region, mock_trainml):
+        """Test wait_for succeeds when status matches."""
+        region._status = "new"
+        api_response_new = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            status="new",
+        )
+        api_response_healthy = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            status="healthy",
+        )
+        mock_trainml._query = AsyncMock(
+            side_effect=[api_response_new, api_response_healthy]
+        )
+        with patch("trainml.cloudbender.regions.asyncio.sleep", new_callable=AsyncMock):
+            result = await region.wait_for("healthy", timeout=10)
+        assert result == region
+        assert region.status == "healthy"
+
+    @mark.asyncio
+    async def test_region_wait_for_archived_404(self, region, mock_trainml):
+        """Test wait_for handles 404 for archived status."""
+        region._status = "healthy"
+        api_error = ApiError(404, {"errorMessage": "Not found"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.regions.asyncio.sleep", new_callable=AsyncMock):
+            await region.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_region_wait_for_error_status(self, region, mock_trainml):
+        """Test wait_for raises error for errored/failed status."""
+        region._status = "new"
+        api_response_errored = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            status="errored",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_errored)
+        with patch("trainml.cloudbender.regions.asyncio.sleep", new_callable=AsyncMock):
+            with raises(specimen.RegionError):
+                await region.wait_for("healthy", timeout=10)
+
+    @mark.asyncio
+    async def test_region_wait_for_timeout(self, region, mock_trainml):
+        """Test wait_for raises timeout exception."""
+        region._status = "new"
+        api_response_new = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            status="new",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_new)
+        with patch("trainml.cloudbender.regions.asyncio.sleep", new_callable=AsyncMock):
+            with raises(TrainMLException) as exc_info:
+                await region.wait_for("healthy", timeout=0.1)
+        assert "Timeout waiting for" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_region_wait_for_api_error_non_404(self, region, mock_trainml):
+        """Test wait_for raises ApiError when not 404 for archived (line 152)."""
+        region._status = "healthy"
+        api_error = ApiError(500, {"errorMessage": "Server Error"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.regions.asyncio.sleep", new_callable=AsyncMock):
+            with raises(ApiError):
+                await region.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_region_wait_for_failed_status(self, region, mock_trainml):
+        """Test wait_for raises error for failed status."""
+        region._status = "new"
+        api_response_failed = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            status="failed",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_failed)
+        with patch("trainml.cloudbender.regions.asyncio.sleep", new_callable=AsyncMock):
+            with raises(specimen.RegionError):
+                await region.wait_for("healthy", timeout=10)

tests/unit/cloudbender/test_services_unit.py

@@ -165,3 +165,144 @@ class serviceTests:
         )
         assert service.id == "service-id-1"
         assert response.id == "service-id-1"
+
+    def test_service_status_property(self, service):
+        """Test service status property."""
+        service._status = "active"
+        assert service.status == "active"
+
+    def test_service_port_property(self, service):
+        """Test service port property."""
+        service._port = "443"
+        assert service.port == "443"
+
+    @mark.asyncio
+    async def test_service_wait_for_already_at_status(self, service):
+        """Test wait_for returns immediately if already at target status."""
+        service._status = "active"
+        result = await service.wait_for("active")
+        assert result is None
+
+    @mark.asyncio
+    async def test_service_wait_for_invalid_status(self, service):
+        """Test wait_for raises error for invalid status."""
+        with raises(SpecificationError) as exc_info:
+            await service.wait_for("invalid_status")
+        assert "Invalid wait_for status" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_service_wait_for_timeout_validation(self, service):
+        """Test wait_for validates timeout."""
+        with raises(SpecificationError) as exc_info:
+            await service.wait_for("active", timeout=25 * 60 * 60)
+        assert "timeout must be less than" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_service_wait_for_success(self, service, mock_trainml):
+        """Test wait_for succeeds when status matches."""
+        service._status = "new"
+        api_response_new = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            service_id="x",
+            status="new",
+        )
+        api_response_active = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            service_id="x",
+            status="active",
+        )
+        mock_trainml._query = AsyncMock(
+            side_effect=[api_response_new, api_response_active]
+        )
+        with patch("trainml.cloudbender.services.asyncio.sleep", new_callable=AsyncMock):
+            result = await service.wait_for("active", timeout=10)
+        assert result == service
+        assert service.status == "active"
+
+    @mark.asyncio
+    async def test_service_wait_for_archived_404(self, service, mock_trainml):
+        """Test wait_for handles 404 for archived status."""
+        service._status = "active"
+        api_error = ApiError(404, {"errorMessage": "Not found"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.services.asyncio.sleep", new_callable=AsyncMock):
+            await service.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_service_wait_for_timeout(self, service, mock_trainml):
+        """Test wait_for raises timeout exception."""
+        service._status = "new"
+        api_response_new = dict(
+            provider_uuid="1",
+            region_uuid="a",
+            service_id="x",
+            status="new",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response_new)
+        with patch("trainml.cloudbender.services.asyncio.sleep", new_callable=AsyncMock):
+            with raises(TrainMLException) as exc_info:
+                await service.wait_for("active", timeout=0.1)
+        assert "Timeout waiting for" in str(exc_info.value.message)
+
+    @mark.asyncio
+    async def test_service_wait_for_api_error_non_404(self, service, mock_trainml):
+        """Test wait_for raises ApiError when not 404 for archived (line 181)."""
+        service._status = "active"
+        api_error = ApiError(500, {"errorMessage": "Server Error"})
+        mock_trainml._query = AsyncMock(side_effect=api_error)
+        with patch("trainml.cloudbender.services.asyncio.sleep", new_callable=AsyncMock):
+            with raises(ApiError):
+                await service.wait_for("archived", timeout=10)
+
+    @mark.asyncio
+    async def test_service_generate_certificate(self, service, mock_trainml):
+        """Test generate_certificate method."""
+        api_response = {
+            "provider_uuid": "1",
+            "region_uuid": "a",
+            "service_id": "x",
+            "certificate": "cert-data",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await service.generate_certificate()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/service/x/certificate",
+            "POST",
+            {},
+            dict(algorithm="ed25519"),
+        )
+        assert result == service
+
+    @mark.asyncio
+    async def test_service_generate_certificate_custom_algorithm(self, service, mock_trainml):
+        """Test generate_certificate with custom algorithm."""
+        api_response = {
+            "provider_uuid": "1",
+            "region_uuid": "a",
+            "service_id": "x",
+            "certificate": "cert-data",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await service.generate_certificate(algorithm="rsa")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/service/x/certificate",
+            "POST",
+            {},
+            dict(algorithm="rsa"),
+        )
+
+    @mark.asyncio
+    async def test_service_sign_client_certificate(self, service, mock_trainml):
+        """Test sign_client_certificate method."""
+        api_response = {"certificate": "signed-cert-data"}
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await service.sign_client_certificate("csr-data")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/service/x/certificate/sign",
+            "POST",
+            {},
+            dict(csr="csr-data"),
+        )
+        assert result == api_response

tests/unit/conftest.py

@@ -4,7 +4,7 @@ from pytest import fixture, mark
 from unittest.mock import Mock, AsyncMock, patch, create_autospec
 
 from trainml.trainml import TrainML
-from trainml.auth import Auth
+from trainml.utils.auth import Auth
 from trainml.datasets import Dataset, Datasets
 from trainml.checkpoints import Checkpoint, Checkpoints
 from trainml.volumes import Volume, Volumes
@@ -12,14 +12,16 @@ from trainml.models import Model, Models
 from trainml.gpu_types import GpuType, GpuTypes
 from trainml.environments import Environment, Environments
 from trainml.jobs import Job, Jobs
-from trainml.connections import Connections
 from trainml.projects import (
     Projects,
     Project,
 )
 from trainml.projects.datastores import ProjectDatastores, ProjectDatastore
 from trainml.projects.services import ProjectServices, ProjectService
-from trainml.projects.data_connectors import ProjectDataConnectors, ProjectDataConnector
+from trainml.projects.data_connectors import (
+    ProjectDataConnectors,
+    ProjectDataConnector,
+)
 from trainml.projects.credentials import ProjectCredentials, ProjectCredential
 from trainml.projects.secrets import ProjectSecrets, ProjectSecret
 
@@ -1130,12 +1132,13 @@ def mock_trainml(
     trainml.gpu_types = create_autospec(GpuTypes)
     trainml.environments = create_autospec(Environments)
     trainml.jobs = create_autospec(Jobs)
-    trainml.connections = create_autospec(Connections)
     trainml.projects = create_autospec(Projects)
     trainml.datasets.list = AsyncMock(return_value=mock_my_datasets)
     trainml.datasets.list_public = AsyncMock(return_value=mock_public_datasets)
     trainml.checkpoints.list = AsyncMock(return_value=mock_my_checkpoints)
-    trainml.checkpoints.list_public = AsyncMock(return_value=mock_public_checkpoints)
+    trainml.checkpoints.list_public = AsyncMock(
+        return_value=mock_public_checkpoints
+    )
     trainml.models.list = AsyncMock(return_value=mock_models)
     trainml.volumes.list = AsyncMock(return_value=mock_my_volumes)
     trainml.gpu_types.list = AsyncMock(return_value=mock_gpu_types)
@@ -1143,17 +1146,25 @@ def mock_trainml(
     trainml.jobs.list = AsyncMock(return_value=mock_jobs)
     trainml.projects.list = AsyncMock(return_value=mock_projects)
     trainml.projects.datastores = create_autospec(ProjectDatastores)
-    trainml.projects.datastores.list = AsyncMock(return_value=mock_project_datastores)
+    trainml.projects.datastores.list = AsyncMock(
+        return_value=mock_project_datastores
+    )
     trainml.projects.services = create_autospec(ProjectServices)
-    trainml.projects.services.list = AsyncMock(return_value=mock_project_services)
+    trainml.projects.services.list = AsyncMock(
+        return_value=mock_project_services
+    )
     trainml.projects.data_connectors = create_autospec(ProjectDataConnectors)
     trainml.projects.data_connectors.list = AsyncMock(
         return_value=mock_project_data_connectors
     )
     trainml.projects.credentials = create_autospec(ProjectCredentials)
-    trainml.projects.credentials.list = AsyncMock(return_value=mock_project_credentials)
+    trainml.projects.credentials.list = AsyncMock(
+        return_value=mock_project_credentials
+    )
     trainml.projects.secrets = create_autospec(ProjectSecrets)
-    trainml.projects.secrets.list = AsyncMock(return_value=mock_project_secrets)
+    trainml.projects.secrets.list = AsyncMock(
+        return_value=mock_project_secrets
+    )
 
     trainml.cloudbender = create_autospec(Cloudbender)
 
@@ -1166,7 +1177,9 @@ def mock_trainml(
     trainml.cloudbender.devices = create_autospec(Devices)
     trainml.cloudbender.devices.list = AsyncMock(return_value=mock_devices)
     trainml.cloudbender.datastores = create_autospec(Datastores)
-    trainml.cloudbender.datastores.list = AsyncMock(return_value=mock_datastores)
+    trainml.cloudbender.datastores.list = AsyncMock(
+        return_value=mock_datastores
+    )
     trainml.cloudbender.services = create_autospec(Services)
     trainml.cloudbender.services.list = AsyncMock(return_value=mock_services)
     trainml.cloudbender.data_connectors = create_autospec(DataConnectors)

tests/unit/projects/test_project_data_connectors_unit.py

@@ -33,6 +33,25 @@ def project_data_connector(mock_trainml):
 
 
 class ProjectDataConnectorsTests:
+    @mark.asyncio
+    async def test_project_data_connectors_get(
+        self, project_data_connectors, mock_trainml
+    ):
+        """Test get method (lines 11-14)."""
+        api_response = {
+            "project_uuid": "proj-id-1",
+            "region_uuid": "reg-id-1",
+            "id": "connector-id-1",
+            "type": "custom",
+            "name": "On-Prem Connection A",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await project_data_connectors.get("connector-id-1", param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/data_connectors/connector-id-1", "GET", dict(param1="value1")
+        )
+        assert result.id == "connector-id-1"
+
     @mark.asyncio
     async def test_project_data_connectors_refresh(
         self, project_data_connectors, mock_trainml
@@ -100,3 +119,23 @@ class ProjectDataConnectorTests:
         empty_project_data_connector = specimen.ProjectDataConnector(mock_trainml)
         assert bool(project_data_connector)
         assert not bool(empty_project_data_connector)
+
+    @mark.asyncio
+    async def test_project_data_connector_enable(self, project_data_connector, mock_trainml):
+        """Test enable method (line 72)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_data_connector.enable()
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/data_connectors/ds-id-1/enable", "PATCH"
+        )
+
+    @mark.asyncio
+    async def test_project_data_connector_disable(self, project_data_connector, mock_trainml):
+        """Test disable method (line 77)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_data_connector.disable()
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/data_connectors/ds-id-1/disable", "PATCH"
+        )

tests/unit/projects/test_project_datastores_unit.py

@@ -33,6 +33,23 @@ def project_datastore(mock_trainml):
 
 
 class ProjectDatastoresTests:
+    @mark.asyncio
+    async def test_project_datastores_get(self, project_datastores, mock_trainml):
+        """Test get method (lines 11-14)."""
+        api_response = {
+            "project_uuid": "proj-id-1",
+            "region_uuid": "reg-id-1",
+            "id": "store-id-1",
+            "type": "nfs",
+            "name": "On-prem NFS",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await project_datastores.get("store-id-1", param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/datastores/store-id-1", "GET", dict(param1="value1")
+        )
+        assert result.id == "store-id-1"
+
     @mark.asyncio
     async def test_project_datastores_refresh(self, project_datastores, mock_trainml):
         api_response = dict()
@@ -94,3 +111,23 @@ class ProjectDatastoreTests:
         empty_project_datastore = specimen.ProjectDatastore(mock_trainml)
         assert bool(project_datastore)
         assert not bool(empty_project_datastore)
+
+    @mark.asyncio
+    async def test_project_datastore_enable(self, project_datastore, mock_trainml):
+        """Test enable method (line 67)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_datastore.enable()
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/datastores/ds-id-1/enable", "PATCH"
+        )
+
+    @mark.asyncio
+    async def test_project_datastore_disable(self, project_datastore, mock_trainml):
+        """Test disable method (line 72)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_datastore.disable()
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/datastores/ds-id-1/disable", "PATCH"
+        )