trainml 0.5.17__py3-none-any.whl → 1.0.1__py3-none-any.whl

tests/unit/projects/test_project_members_unit.py

@@ -37,6 +37,42 @@ def project_member(mock_trainml):
 
 
 class ProjectMembersTests:
+    @mark.asyncio
+    async def test_project_members_add(self, project_members, mock_trainml):
+        """Test add method (lines 18-31)."""
+        api_response = {
+            "project_uuid": "proj-id-1",
+            "email": "newuser@gmail.com",
+            "owner": False,
+            "job": "all",
+            "dataset": "read",
+            "model": "all",
+            "checkpoint": "read",
+            "volume": "all",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await project_members.add(
+            email="newuser@gmail.com",
+            job="all",
+            dataset="read",
+            model="all",
+            checkpoint="read",
+            volume="all",
+            param1="value1",
+        )
+        expected_payload = dict(
+            email="newuser@gmail.com",
+            job="all",
+            dataset="read",
+            model="all",
+            checkpoint="read",
+            volume="all",
+        )
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/access", "POST", dict(param1="value1"), expected_payload
+        )
+        assert result.email == "newuser@gmail.com"
+
     @mark.asyncio
     async def test_project_members_list(self, project_members, mock_trainml):
         api_response = [
@@ -72,6 +108,16 @@ class ProjectMembersTests:
         )
         assert len(resp) == 2
 
+    @mark.asyncio
+    async def test_project_members_remove(self, project_members, mock_trainml):
+        """Test remove method (line 35)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_members.remove("user@gmail.com", param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/access", "DELETE", dict(param1="value1", email="user@gmail.com")
+        )
+
 
 class ProjectMemberTests:
     def test_project_member_properties(self, project_member):

tests/unit/projects/test_project_services_unit.py

@@ -34,6 +34,24 @@ def project_service(mock_trainml):
 
 
 class ProjectServicesTests:
+    @mark.asyncio
+    async def test_project_services_get(self, project_services, mock_trainml):
+        """Test get method (lines 11-14)."""
+        api_response = {
+            "project_uuid": "proj-id-1",
+            "region_uuid": "reg-id-1",
+            "id": "res-id-1",
+            "type": "port",
+            "name": "On-Prem Service A",
+            "hostname": "service-a.local",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await project_services.get("res-id-1", param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/services/res-id-1", "GET", dict(param1="value1")
+        )
+        assert result.id == "res-id-1"
+
     @mark.asyncio
     async def test_project_services_refresh(self, project_services, mock_trainml):
         api_response = dict()
@@ -100,3 +118,50 @@ class ProjectServiceTests:
         empty_project_service = specimen.ProjectService(mock_trainml)
         assert bool(project_service)
         assert not bool(empty_project_service)
+
+    @mark.asyncio
+    async def test_project_service_enable(self, project_service, mock_trainml):
+        """Test enable method (line 72)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_service.enable()
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/services/res-id-1/enable", "PATCH"
+        )
+
+    @mark.asyncio
+    async def test_project_service_disable(self, project_service, mock_trainml):
+        """Test disable method (line 77)."""
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_service.disable()
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/services/res-id-1/disable", "PATCH"
+        )
+
+    @mark.asyncio
+    async def test_project_service_get_service_ca_certificate(self, project_service, mock_trainml):
+        """Test get_service_ca_certificate method (lines 82-87)."""
+        api_response = {"certificate": "ca-cert-data"}
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await project_service.get_service_ca_certificate(param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/services/res-id-1/certificate/ca",
+            "GET",
+            dict(param1="value1"),
+        )
+        assert result == api_response
+
+    @mark.asyncio
+    async def test_project_service_sign_client_certificate(self, project_service, mock_trainml):
+        """Test sign_client_certificate method (lines 90-96)."""
+        api_response = {"certificate": "signed-cert-data"}
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await project_service.sign_client_certificate("csr-data", param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/proj-id-1/services/res-id-1/certificate/sign",
+            "POST",
+            dict(param1="value1"),
+            dict(csr="csr-data"),
+        )
+        assert result == api_response

tests/unit/projects/test_projects_unit.py

@@ -48,6 +48,22 @@ class ProjectsTests:
         await projects.get("1234")
         mock_trainml._query.assert_called_once_with("/project/1234", "GET", dict())
 
+    @mark.asyncio
+    async def test_get_current_project(self, projects, mock_trainml):
+        """Test get_current method (lines 20-23)."""
+        api_response = {
+            "id": "project-id-1",
+            "name": "current project",
+            "owner": True,
+        }
+        mock_trainml.project = "project-id-1"
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await projects.get_current(param1="value1")
+        mock_trainml._query.assert_called_once_with(
+            "/project/project-id-1", "GET", dict(param1="value1")
+        )
+        assert result.id == "project-id-1"
+
     @mark.asyncio
     async def test_list_projects(
         self,

tests/unit/test_auth_unit.py

@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, patch, mock_open, MagicMock
 from pytest import mark, fixture, raises
 from aiohttp import WSMessage, WSMsgType
 
-import trainml.auth as specimen
+import trainml.utils.auth as specimen
 
 pytestmark = [mark.sdk, mark.unit]
 
@@ -21,7 +21,22 @@ pytestmark = [mark.sdk, mark.unit]
         "TRAINML_POOL_ID": "pool_id",
     },
 )
-def test_auth_from_envs():
+@patch("trainml.utils.auth.boto3.client")
+@patch("trainml.utils.auth.requests.get")
+@patch("builtins.open", side_effect=FileNotFoundError)
+def test_auth_from_envs(mock_open, mock_requests_get, mock_boto3_client):
+    # Mock the auth config request
+    mock_response = MagicMock()
+    mock_response.json.return_value = {
+        "region": "us-east-1",
+        "userPoolSDKClientId": "default_client_id",
+        "userPoolId": "default_pool_id",
+    }
+    mock_requests_get.return_value = mock_response
+
+    # Mock boto3 client
+    mock_boto3_client.return_value = MagicMock()
+
     auth = specimen.Auth(config_dir=os.path.expanduser("~/.trainml"))
     assert auth.__dict__.get("username") == "user-id"
     assert auth.__dict__.get("password") == "key"

tests/unit/test_checkpoints_unit.py

@@ -44,7 +44,9 @@ class CheckpointsTests:
         api_response = dict()
         mock_trainml._query = AsyncMock(return_value=api_response)
         await checkpoints.get("1234")
-        mock_trainml._query.assert_called_once_with("/checkpoint/1234", "GET", dict())
+        mock_trainml._query.assert_called_once_with(
+            "/checkpoint/1234", "GET", dict()
+        )
 
     @mark.asyncio
     async def test_list_checkpoints(
@@ -55,7 +57,30 @@ class CheckpointsTests:
         api_response = dict()
         mock_trainml._query = AsyncMock(return_value=api_response)
         await checkpoints.list()
-        mock_trainml._query.assert_called_once_with("/checkpoint", "GET", dict())
+        mock_trainml._query.assert_called_once_with(
+            "/checkpoint", "GET", dict()
+        )
+
+    @mark.asyncio
+    async def test_list_public_checkpoints(
+        self,
+        checkpoints,
+        mock_trainml,
+    ):
+        api_response = [
+            dict(
+                checkpoint_uuid="1",
+                name="public checkpoint",
+                status="ready",
+            )
+        ]
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await checkpoints.list_public()
+        mock_trainml._query.assert_called_once_with(
+            "/checkpoint/public", "GET", dict()
+        )
+        assert len(result) == 1
+        assert isinstance(result[0], specimen.Checkpoint)
 
     @mark.asyncio
     async def test_remove_checkpoint(
@@ -133,9 +158,7 @@ class CheckpointTests:
 
     @mark.asyncio
     async def test_checkpoint_get_log_url(self, checkpoint, mock_trainml):
-        api_response = (
-            "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/logs/first_one.zip"
-        )
+        api_response = "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/logs/first_one.zip"
         mock_trainml._query = AsyncMock(return_value=api_response)
         response = await checkpoint.get_log_url()
         mock_trainml._query.assert_called_once_with(
@@ -160,81 +183,155 @@ class CheckpointTests:
         assert response == api_response
 
     @mark.asyncio
-    async def test_checkpoint_get_connection_utility_url(
-        self, checkpoint, mock_trainml
-    ):
-        api_response = (
-            "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/vpn/first_one.zip"
+    async def test_checkpoint_connect_downloading_status(self, mock_trainml):
+        checkpoint = specimen.Checkpoint(
+            mock_trainml,
+            checkpoint_uuid="1",
+            project_uuid="proj-id-1",
+            name="test checkpoint",
+            status="downloading",
+            auth_token="test-token",
+            hostname="example.com",
+            source_uri="/path/to/source",
         )
-        mock_trainml._query = AsyncMock(return_value=api_response)
-        response = await checkpoint.get_connection_utility_url()
-        mock_trainml._query.assert_called_once_with(
-            "/checkpoint/1/download", "GET", dict(project_uuid="proj-id-1")
+
+        with patch(
+            "trainml.checkpoints.Checkpoint.refresh", new_callable=AsyncMock
+        ) as mock_refresh:
+            with patch(
+                "trainml.checkpoints.upload", new_callable=AsyncMock
+            ) as mock_upload:
+                await checkpoint.connect()
+                mock_refresh.assert_called_once()
+                mock_upload.assert_called_once_with(
+                    "example.com", "test-token", "/path/to/source"
+                )
+
+    @mark.asyncio
+    async def test_checkpoint_connect_exporting_status(
+        self, mock_trainml, tmp_path
+    ):
+        output_dir = str(tmp_path / "output")
+        checkpoint = specimen.Checkpoint(
+            mock_trainml,
+            checkpoint_uuid="1",
+            project_uuid="proj-id-1",
+            name="test checkpoint",
+            status="exporting",
+            auth_token="test-token",
+            hostname="example.com",
+            output_uri=output_dir,
         )
-        assert response == api_response
 
-    def test_checkpoint_get_connection_details_no_vpn(self, checkpoint):
-        details = checkpoint.get_connection_details()
-        expected_details = dict()
-        assert details == expected_details
+        with patch(
+            "trainml.checkpoints.Checkpoint.refresh", new_callable=AsyncMock
+        ) as mock_refresh:
+            with patch(
+                "trainml.checkpoints.download", new_callable=AsyncMock
+            ) as mock_download:
+                await checkpoint.connect()
+                mock_refresh.assert_called_once()
+                mock_download.assert_called_once_with(
+                    "example.com", "test-token", output_dir
+                )
 
-    def test_checkpoint_get_connection_details_local_data(self, mock_trainml):
+    @mark.asyncio
+    async def test_checkpoint_connect_new_status_waits_for_downloading(
+        self, mock_trainml
+    ):
         checkpoint = specimen.Checkpoint(
             mock_trainml,
             checkpoint_uuid="1",
-            project_uuid="a",
-            name="first one",
+            project_uuid="proj-id-1",
+            name="test checkpoint",
             status="new",
-            size=100000,
-            createdAt="2020-12-31T23:59:59.000Z",
-            source_type="local",
-            source_uri="~/tensorflow-example",
-            vpn={
-                "status": "new",
-                "cidr": "10.106.171.0/24",
-                "client": {
-                    "port": "36017",
-                    "id": "cus-id-1",
-                    "address": "10.106.171.253",
-                    "ssh_port": 46600,
-                },
-                "net_prefix_type_id": 1,
-            },
-        )
-        details = checkpoint.get_connection_details()
-        expected_details = dict(
-            project_uuid="a",
-            entity_type="checkpoint",
-            cidr="10.106.171.0/24",
-            ssh_port=46600,
-            input_path="~/tensorflow-example",
-            output_path=None,
-        )
-        assert details == expected_details
+        )
+
+        with patch(
+            "trainml.checkpoints.Checkpoint.wait_for", new_callable=AsyncMock
+        ) as mock_wait:
+            with patch(
+                "trainml.checkpoints.Checkpoint.refresh",
+                new_callable=AsyncMock,
+            ) as mock_refresh:
+                # After refresh, status becomes downloading
+                def update_status():
+                    checkpoint._status = "downloading"
+                    checkpoint._checkpoint.update(
+                        {
+                            "auth_token": "test-token",
+                            "hostname": "example.com",
+                            "source_uri": "/path/to/source",
+                        }
+                    )
+
+                mock_refresh.side_effect = update_status
+
+                with patch(
+                    "trainml.checkpoints.upload", new_callable=AsyncMock
+                ) as mock_upload:
+                    await checkpoint.connect()
+                    mock_wait.assert_called_once_with("downloading")
+                    mock_refresh.assert_called_once()
+                    mock_upload.assert_called_once()
 
     @mark.asyncio
-    async def test_checkpoint_connect(self, checkpoint, mock_trainml):
+    async def test_checkpoint_connect_invalid_status(self, mock_trainml):
+        checkpoint = specimen.Checkpoint(
+            mock_trainml,
+            checkpoint_uuid="1",
+            project_uuid="proj-id-1",
+            name="test checkpoint",
+            status="ready",
+        )
+
+        with raises(
+            SpecificationError,
+            match="You can only connect to downloading or exporting checkpoints",
+        ):
+            await checkpoint.connect()
+
+    @mark.asyncio
+    async def test_checkpoint_connect_missing_properties_downloading(
+        self, mock_trainml
+    ):
+        checkpoint = specimen.Checkpoint(
+            mock_trainml,
+            checkpoint_uuid="1",
+            project_uuid="proj-id-1",
+            name="test checkpoint",
+            status="downloading",
+        )
+
         with patch(
-            "trainml.checkpoints.Connection",
-            autospec=True,
-        ) as mock_connection:
-            connection = mock_connection.return_value
-            connection.status = "connected"
-            resp = await checkpoint.connect()
-            connection.start.assert_called_once()
-            assert resp == "connected"
+            "trainml.checkpoints.Checkpoint.refresh", new_callable=AsyncMock
+        ):
+            with raises(
+                SpecificationError,
+                match="missing required connection properties",
+            ):
+                await checkpoint.connect()
 
     @mark.asyncio
-    async def test_checkpoint_disconnect(self, checkpoint, mock_trainml):
+    async def test_checkpoint_connect_missing_properties_exporting(
+        self, mock_trainml
+    ):
+        checkpoint = specimen.Checkpoint(
+            mock_trainml,
+            checkpoint_uuid="1",
+            project_uuid="proj-id-1",
+            name="test checkpoint",
+            status="exporting",
+        )
+
         with patch(
-            "trainml.checkpoints.Connection",
-            autospec=True,
-        ) as mock_connection:
-            connection = mock_connection.return_value
-            connection.status = "removed"
-            resp = await checkpoint.disconnect()
-            connection.stop.assert_called_once()
-            assert resp == "removed"
+            "trainml.checkpoints.Checkpoint.refresh", new_callable=AsyncMock
+        ):
+            with raises(
+                SpecificationError,
+                match="missing required connection properties",
+            ):
+                await checkpoint.connect()
 
     @mark.asyncio
     async def test_checkpoint_remove(self, checkpoint, mock_trainml):
@@ -340,7 +437,9 @@ class CheckpointTests:
         assert response.id == "data-id-1"
 
     @mark.asyncio
-    async def test_checkpoint_wait_for_successful(self, checkpoint, mock_trainml):
+    async def test_checkpoint_wait_for_successful(
+        self, checkpoint, mock_trainml
+    ):
         api_response = {
             "customer_uuid": "cus-id-1",
             "checkpoint_uuid": "data-id-1",
@@ -373,7 +472,9 @@ class CheckpointTests:
         mock_trainml._query.assert_not_called()
 
     @mark.asyncio
-    async def test_checkpoint_wait_for_incorrect_status(self, checkpoint, mock_trainml):
+    async def test_checkpoint_wait_for_incorrect_status(
+        self, checkpoint, mock_trainml
+    ):
         api_response = None
         mock_trainml._query = AsyncMock(return_value=api_response)
         with raises(SpecificationError):
@@ -381,7 +482,9 @@ class CheckpointTests:
         mock_trainml._query.assert_not_called()
 
     @mark.asyncio
-    async def test_checkpoint_wait_for_with_delay(self, checkpoint, mock_trainml):
+    async def test_checkpoint_wait_for_with_delay(
+        self, checkpoint, mock_trainml
+    ):
         api_response_initial = dict(
             checkpoint_uuid="1",
             name="first one",
@@ -437,7 +540,9 @@ class CheckpointTests:
         self, checkpoint, mock_trainml
     ):
         mock_trainml._query = AsyncMock(
-            side_effect=ApiError(404, dict(errorMessage="Checkpoint Not Found"))
+            side_effect=ApiError(
+                404, dict(errorMessage="Checkpoint Not Found")
+            )
         )
         await checkpoint.wait_for("archived")
         mock_trainml._query.assert_called()
@@ -447,8 +552,88 @@ class CheckpointTests:
         self, checkpoint, mock_trainml
     ):
         mock_trainml._query = AsyncMock(
-            side_effect=ApiError(404, dict(errorMessage="Checkpoint Not Found"))
+            side_effect=ApiError(
+                404, dict(errorMessage="Checkpoint Not Found")
+            )
         )
         with raises(ApiError):
             await checkpoint.wait_for("ready")
         mock_trainml._query.assert_called()
+
+    @mark.asyncio
+    async def test_checkpoint_rename(self, checkpoint, mock_trainml):
+        api_response = dict(
+            checkpoint_uuid="1",
+            name="renamed checkpoint",
+            project_uuid="proj-id-1",
+            status="ready",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await checkpoint.rename("renamed checkpoint")
+        mock_trainml._query.assert_called_once_with(
+            "/checkpoint/1",
+            "PATCH",
+            dict(project_uuid="proj-id-1"),
+            dict(name="renamed checkpoint"),
+        )
+        assert result == checkpoint
+        assert checkpoint.name == "renamed checkpoint"
+
+    @mark.asyncio
+    async def test_checkpoint_export(self, checkpoint, mock_trainml):
+        api_response = dict(
+            checkpoint_uuid="1",
+            name="first one",
+            project_uuid="proj-id-1",
+            status="exporting",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await checkpoint.export("aws", "s3://bucket/path", dict(key="value"))
+        mock_trainml._query.assert_called_once_with(
+            "/checkpoint/1/export",
+            "POST",
+            dict(project_uuid="proj-id-1"),
+            dict(
+                output_type="aws",
+                output_uri="s3://bucket/path",
+                output_options=dict(key="value"),
+            ),
+        )
+        assert result == checkpoint
+        assert checkpoint.status == "exporting"
+
+    @mark.asyncio
+    async def test_checkpoint_export_default_options(self, checkpoint, mock_trainml):
+        api_response = dict(
+            checkpoint_uuid="1",
+            name="first one",
+            project_uuid="proj-id-1",
+            status="exporting",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await checkpoint.export("aws", "s3://bucket/path")
+        mock_trainml._query.assert_called_once_with(
+            "/checkpoint/1/export",
+            "POST",
+            dict(project_uuid="proj-id-1"),
+            dict(
+                output_type="aws",
+                output_uri="s3://bucket/path",
+                output_options=dict(),
+            ),
+        )
+        assert result == checkpoint
+
+    @mark.asyncio
+    async def test_checkpoint_wait_for_timeout_validation(
+        self, checkpoint, mock_trainml
+    ):
+        with raises(SpecificationError) as exc_info:
+            await checkpoint.wait_for("ready", timeout=25 * 60 * 60)  # > 24 hours
+        assert "timeout" in str(exc_info.value.attribute).lower()
+        assert "less than" in str(exc_info.value.message).lower()
+
+    def test_checkpoint_billed_size_property(self, checkpoint, mock_trainml):
+        """Test billed_size property access."""
+        checkpoint._billed_size = 50000
+        assert checkpoint.billed_size == 50000