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

tests/unit/test_datasets_unit.py

@@ -142,7 +142,9 @@ class DatasetTests:
     def test_dataset_repr(self, dataset):
         string = repr(dataset)
         regex = (
-            r"^Dataset\( trainml , \*\*{.*'dataset_uuid': '" + dataset.id + r"'.*}\)$"
+            r"^Dataset\( trainml , \*\*{.*'dataset_uuid': '"
+            + dataset.id
+            + r"'.*}\)$"
         )
         assert isinstance(string, str)
         assert re.match(regex, string)
@@ -154,9 +156,7 @@ class DatasetTests:
 
     @mark.asyncio
     async def test_dataset_get_log_url(self, dataset, 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 dataset.get_log_url()
         mock_trainml._query.assert_called_once_with(
@@ -181,78 +181,73 @@ class DatasetTests:
         assert response == api_response
 
     @mark.asyncio
-    async def test_dataset_get_connection_utility_url(self, dataset, mock_trainml):
-        api_response = (
-            "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/vpn/first_one.zip"
-        )
-        mock_trainml._query = AsyncMock(return_value=api_response)
-        response = await dataset.get_connection_utility_url()
-        mock_trainml._query.assert_called_once_with(
-            "/dataset/1/download", "GET", dict(project_uuid="proj-id-1")
+    async def test_dataset_connect_downloading_status(self, mock_trainml):
+        dataset = specimen.Dataset(
+            mock_trainml,
+            dataset_uuid="1",
+            project_uuid="proj-id-1",
+            name="test dataset",
+            status="downloading",
+            auth_token="test-token",
+            hostname="example.com",
+            source_uri="/path/to/source",
         )
-        assert response == api_response
 
-    def test_dataset_get_connection_details_no_vpn(self, dataset):
-        details = dataset.get_connection_details()
-        expected_details = dict()
-        assert details == expected_details
+        with patch(
+            "trainml.datasets.Dataset.refresh", new_callable=AsyncMock
+        ) as mock_refresh:
+            with patch(
+                "trainml.datasets.upload", new_callable=AsyncMock
+            ) as mock_upload:
+                await dataset.connect()
+                mock_refresh.assert_called_once()
+                mock_upload.assert_called_once_with(
+                    "example.com", "test-token", "/path/to/source"
+                )
 
-    def test_dataset_get_connection_details_local_data(self, mock_trainml):
+    @mark.asyncio
+    async def test_dataset_connect_exporting_status(
+        self, mock_trainml, tmp_path
+    ):
+        output_dir = str(tmp_path / "output")
         dataset = specimen.Dataset(
             mock_trainml,
             dataset_uuid="1",
             project_uuid="proj-id-1",
-            name="first one",
-            status="new",
-            size=100000,
-            createdAt="2020-12-31T23:59:59.000Z",
-            source_type="local",
-            source_uri="~/tensorflow-example/data",
-            vpn={
-                "status": "new",
-                "cidr": "10.106.171.0/24",
-                "client": {
-                    "port": "36017",
-                    "id": "cus-id-1",
-                    "address": "10.106.171.253",
-                    "ssh_port": 46600,
-                },
-            },
-        )
-        details = dataset.get_connection_details()
-        expected_details = dict(
-            project_uuid="proj-id-1",
-            entity_type="dataset",
-            cidr="10.106.171.0/24",
-            ssh_port=46600,
-            input_path="~/tensorflow-example/data",
-            output_path=None,
+            name="test dataset",
+            status="exporting",
+            auth_token="test-token",
+            hostname="example.com",
+            output_uri=output_dir,
         )
-        assert details == expected_details
 
-    @mark.asyncio
-    async def test_dataset_connect(self, dataset, mock_trainml):
         with patch(
-            "trainml.datasets.Connection",
-            autospec=True,
-        ) as mock_connection:
-            connection = mock_connection.return_value
-            connection.status = "connected"
-            resp = await dataset.connect()
-            connection.start.assert_called_once()
-            assert resp == "connected"
+            "trainml.datasets.Dataset.refresh", new_callable=AsyncMock
+        ) as mock_refresh:
+            with patch(
+                "trainml.datasets.download", new_callable=AsyncMock
+            ) as mock_download:
+                await dataset.connect()
+                mock_refresh.assert_called_once()
+                mock_download.assert_called_once_with(
+                    "example.com", "test-token", output_dir
+                )
 
     @mark.asyncio
-    async def test_dataset_disconnect(self, dataset, mock_trainml):
-        with patch(
-            "trainml.datasets.Connection",
-            autospec=True,
-        ) as mock_connection:
-            connection = mock_connection.return_value
-            connection.status = "removed"
-            resp = await dataset.disconnect()
-            connection.stop.assert_called_once()
-            assert resp == "removed"
+    async def test_dataset_connect_invalid_status(self, mock_trainml):
+        dataset = specimen.Dataset(
+            mock_trainml,
+            dataset_uuid="1",
+            project_uuid="proj-id-1",
+            name="test dataset",
+            status="ready",
+        )
+
+        with raises(
+            SpecificationError,
+            match="You can only connect to downloading or exporting datasets",
+        ):
+            await dataset.connect()
 
     @mark.asyncio
     async def test_dataset_remove(self, dataset, mock_trainml):
@@ -391,7 +386,9 @@ class DatasetTests:
         mock_trainml._query.assert_not_called()
 
     @mark.asyncio
-    async def test_dataset_wait_for_incorrect_status(self, dataset, mock_trainml):
+    async def test_dataset_wait_for_incorrect_status(
+        self, dataset, mock_trainml
+    ):
         api_response = None
         mock_trainml._query = AsyncMock(return_value=api_response)
         with raises(SpecificationError):
@@ -436,7 +433,9 @@ class DatasetTests:
         mock_trainml._query.assert_called()
 
     @mark.asyncio
-    async def test_dataset_wait_for_dataset_failed(self, dataset, mock_trainml):
+    async def test_dataset_wait_for_dataset_failed(
+        self, dataset, mock_trainml
+    ):
         api_response = dict(
             dataset_uuid="1",
             name="first one",
@@ -449,7 +448,156 @@ class DatasetTests:
         mock_trainml._query.assert_called()
 
     @mark.asyncio
-    async def test_dataset_wait_for_archived_succeeded(self, dataset, mock_trainml):
+    async def test_dataset_rename(self, dataset, mock_trainml):
+        api_response = dict(
+            dataset_uuid="1",
+            name="renamed dataset",
+            project_uuid="proj-id-1",
+            status="ready",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await dataset.rename("renamed dataset")
+        mock_trainml._query.assert_called_once_with(
+            "/dataset/1",
+            "PATCH",
+            None,
+            dict(name="renamed dataset"),
+        )
+        assert result == dataset
+        assert dataset.name == "renamed dataset"
+
+    @mark.asyncio
+    async def test_dataset_export(self, dataset, mock_trainml):
+        api_response = dict(
+            dataset_uuid="1",
+            name="first one",
+            project_uuid="proj-id-1",
+            status="exporting",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await dataset.export("aws", "s3://bucket/path", dict(key="value"))
+        mock_trainml._query.assert_called_once_with(
+            "/dataset/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 == dataset
+        assert dataset.status == "exporting"
+
+    @mark.asyncio
+    async def test_dataset_export_default_options(self, dataset, mock_trainml):
+        api_response = dict(
+            dataset_uuid="1",
+            name="first one",
+            project_uuid="proj-id-1",
+            status="exporting",
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        result = await dataset.export("aws", "s3://bucket/path")
+        mock_trainml._query.assert_called_once_with(
+            "/dataset/1/export",
+            "POST",
+            dict(project_uuid="proj-id-1"),
+            dict(
+                output_type="aws",
+                output_uri="s3://bucket/path",
+                output_options=dict(),
+            ),
+        )
+        assert result == dataset
+
+    @mark.asyncio
+    async def test_dataset_wait_for_timeout_validation(
+        self, dataset, mock_trainml
+    ):
+        with raises(SpecificationError) as exc_info:
+            await dataset.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()
+
+    @mark.asyncio
+    async def test_dataset_connect_new_status_waits_for_downloading(
+        self, dataset, mock_trainml
+    ):
+        """Test that connect() waits for downloading status when status is 'new'."""
+        dataset._dataset["status"] = "new"
+        dataset._status = "new"
+        api_response_new = dict(
+            dataset_uuid="1",
+            name="first one",
+            status="new",
+        )
+        api_response_downloading = dict(
+            dataset_uuid="1",
+            name="first one",
+            status="downloading",
+            auth_token="token",
+            hostname="host",
+            source_uri="s3://bucket/path",
+        )
+        # wait_for calls refresh multiple times, then connect calls refresh again
+        # We need enough responses for wait_for polling and the final refresh
+        mock_trainml._query = AsyncMock(
+            side_effect=[
+                api_response_new,  # wait_for refresh 1
+                api_response_downloading,  # wait_for refresh 2 (status matches, wait_for returns)
+                api_response_downloading,  # connect refresh
+            ]
+        )
+        with patch("trainml.datasets.upload", new_callable=AsyncMock) as mock_upload:
+            await dataset.connect()
+        # After refresh, status should be downloading
+        assert dataset.status == "downloading"
+        mock_upload.assert_called_once()
+
+    @mark.asyncio
+    async def test_dataset_connect_downloading_missing_properties(
+        self, dataset, mock_trainml
+    ):
+        """Test connect() raises error when downloading status missing properties."""
+        dataset._dataset["status"] = "downloading"
+        api_response = dict(
+            dataset_uuid="1",
+            name="first one",
+            status="downloading",
+            # Missing auth_token, hostname, or source_uri
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        with raises(SpecificationError) as exc_info:
+            await dataset.connect()
+        assert "missing required connection properties" in str(exc_info.value.message).lower()
+
+    @mark.asyncio
+    async def test_dataset_connect_exporting_missing_properties(
+        self, dataset, mock_trainml
+    ):
+        """Test connect() raises error when exporting status missing properties."""
+        dataset._dataset["status"] = "exporting"
+        api_response = dict(
+            dataset_uuid="1",
+            name="first one",
+            status="exporting",
+            # Missing auth_token, hostname, or output_uri
+        )
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        with raises(SpecificationError) as exc_info:
+            await dataset.connect()
+        assert "missing required connection properties" in str(exc_info.value.message).lower()
+
+    def test_dataset_billed_size_property(self, dataset, mock_trainml):
+        """Test billed_size property access."""
+        dataset._billed_size = 50000
+        assert dataset.billed_size == 50000
+
+    @mark.asyncio
+    async def test_dataset_wait_for_archived_succeeded(
+        self, dataset, mock_trainml
+    ):
         mock_trainml._query = AsyncMock(
             side_effect=ApiError(404, dict(errorMessage="Dataset Not Found"))
         )
@@ -457,7 +605,9 @@ class DatasetTests:
         mock_trainml._query.assert_called()
 
     @mark.asyncio
-    async def test_dataset_wait_for_unexpected_api_error(self, dataset, mock_trainml):
+    async def test_dataset_wait_for_unexpected_api_error(
+        self, dataset, mock_trainml
+    ):
         mock_trainml._query = AsyncMock(
             side_effect=ApiError(404, dict(errorMessage="Dataset Not Found"))
         )

tests/unit/test_exceptions.py

@@ -5,22 +5,55 @@ import trainml.exceptions as specimen
 pytestmark = [mark.sdk, mark.unit]
 
 
+def test_trainml_exception():
+    """Test TrainMLException base class."""
+    error = specimen.TrainMLException("test message")
+    assert error.message == "test message"
+    assert repr(error) == "TrainMLException( 'test message')"
+    assert str(error) == "TrainMLException('test message')"
+    
+    # Test with multiple args
+    error2 = specimen.TrainMLException("test", "arg1", "arg2")
+    assert error2.message == "test"
+
+
 def test_api_error():
+    """Test ApiError exception."""
     error = specimen.ApiError(400, dict(errorMessage="test message"))
+    assert error.status == 400
+    assert error.message == "test message"
     assert repr(error) == "ApiError(400, 'test message')"
     assert str(error) == "ApiError(400, 'test message')"
+    
+    # Test with message key instead of errorMessage
+    error2 = specimen.ApiError(404, dict(message="not found"))
+    assert error2.message == "not found"
+    
+    # Test with multiple args
+    error3 = specimen.ApiError(500, dict(errorMessage="server error"), "extra")
+    assert error3.message == "server error"
 
 
 def test_job_error():
+    """Test JobError exception."""
     error = specimen.JobError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
     assert (
         repr(error) == "JobError(failed, {'id': 'id-1', 'status': 'failed'})"
     )
     assert str(error) == "JobError(failed, {'id': 'id-1', 'status': 'failed'})"
+    
+    # Test with string data
+    error2 = specimen.JobError("errored", "error string")
+    assert error2.message == "error string"
 
 
 def test_dataset_error():
+    """Test DatasetError exception."""
     error = specimen.DatasetError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
     assert (
         repr(error)
         == "DatasetError(failed, {'id': 'id-1', 'status': 'failed'})"
@@ -29,3 +62,103 @@ def test_dataset_error():
         str(error)
         == "DatasetError(failed, {'id': 'id-1', 'status': 'failed'})"
     )
+
+
+def test_model_error():
+    """Test ModelError exception."""
+    error = specimen.ModelError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
+    assert (
+        repr(error) == "ModelError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+    assert str(error) == "ModelError(failed, {'id': 'id-1', 'status': 'failed'})"
+
+
+def test_checkpoint_error():
+    """Test CheckpointError exception."""
+    error = specimen.CheckpointError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
+    assert (
+        repr(error)
+        == "CheckpointError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+    assert (
+        str(error)
+        == "CheckpointError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+
+
+def test_volume_error():
+    """Test VolumeError exception."""
+    error = specimen.VolumeError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
+    assert (
+        repr(error) == "VolumeError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+    assert str(error) == "VolumeError(failed, {'id': 'id-1', 'status': 'failed'})"
+
+
+def test_connection_error():
+    """Test ConnectionError exception."""
+    error = specimen.ConnectionError("connection failed")
+    assert error.message == "connection failed"
+    assert repr(error) == "ConnectionError(connection failed)"
+    assert str(error) == "ConnectionError(connection failed)"
+    
+    # Test with multiple args
+    error2 = specimen.ConnectionError("test", "arg1", "arg2")
+    assert error2.message == "test"
+
+
+def test_specification_error():
+    """Test SpecificationError exception."""
+    error = specimen.SpecificationError("attr", "invalid value")
+    assert error.attribute == "attr"
+    assert error.message == "invalid value"
+    assert repr(error) == "SpecificationError(attr, invalid value)"
+    assert str(error) == "SpecificationError(attr, invalid value)"
+    
+    # Test with multiple args
+    error2 = specimen.SpecificationError("attr", "test", "arg1")
+    assert error2.attribute == "attr"
+    assert error2.message == "test"
+
+
+def test_node_error():
+    """Test NodeError exception."""
+    error = specimen.NodeError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
+    assert (
+        repr(error) == "NodeError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+    assert str(error) == "NodeError(failed, {'id': 'id-1', 'status': 'failed'})"
+
+
+def test_provider_error():
+    """Test ProviderError exception."""
+    error = specimen.ProviderError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
+    assert (
+        repr(error)
+        == "ProviderError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+    assert (
+        str(error)
+        == "ProviderError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+
+
+def test_region_error():
+    """Test RegionError exception."""
+    error = specimen.RegionError("failed", dict(id="id-1", status="failed"))
+    assert error.status == "failed"
+    assert error.message == dict(id="id-1", status="failed")
+    assert (
+        repr(error) == "RegionError(failed, {'id': 'id-1', 'status': 'failed'})"
+    )
+    assert str(error) == "RegionError(failed, {'id': 'id-1', 'status': 'failed'})"

tests/unit/test_gpu_types_unit.py

@@ -1,6 +1,6 @@
 import re
 from unittest.mock import AsyncMock
-from pytest import mark, fixture
+from pytest import mark, fixture, raises
 
 import trainml.gpu_types as specimen
 
@@ -39,6 +39,16 @@ class GpuTypesTests:
             f"/project/proj-id-1/gputypes", "GET"
         )
 
+    @mark.asyncio
+    async def test_list_gpu_types_no_project(self, mock_trainml):
+        """Test list raises error when no active project (line 11)."""
+        from trainml.exceptions import TrainMLException
+        gpu_types = specimen.GpuTypes(mock_trainml)
+        mock_trainml.project = None
+        with raises(TrainMLException) as exc_info:
+            await gpu_types.list()
+        assert "Active project not configured" in str(exc_info.value.message)
+
     @mark.asyncio
     async def test_project_refresh_gpu_types(self, gpu_types, mock_trainml):
         api_response = dict()