truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/tests/patch/test_dir_signature.py

@@ -1,4 +1,4 @@
-from truss.patch.dir_signature import directory_content_signature
+from truss.truss_handle.patch.dir_signature import directory_content_signature
 
 
 def test_directory_content_signature(tmp_path):
@@ -12,12 +12,7 @@ def test_directory_content_signature(tmp_path):
 
     content_sign = directory_content_signature(root)
 
-    assert content_sign.keys() == {
-        "dir",
-        "dir/file3",
-        "file1",
-        "file2",
-    }
+    assert content_sign.keys() == {"dir", "dir/file3", "file1", "file2"}
 
 
 def test_directory_content_signature_ignore_patterns(tmp_path):
@@ -40,8 +35,4 @@ def test_directory_content_signature_ignore_patterns(tmp_path):
         root=root, ignore_patterns=["data/*", ".git"]
     )
 
-    assert content_sign.keys() == {
-        "data",
-        "file1",
-        "file2",
-    }
+    assert content_sign.keys() == {"data", "file1", "file2"}

truss/tests/patch/test_hash.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import Callable, List
 
 import pytest
-from truss.patch.hash import (
+from truss.truss_handle.patch.hash import (
     directory_content_hash,
     file_content_hash,
     file_content_hash_str,

truss/tests/patch/test_signature.py

@@ -1,4 +1,4 @@
-from truss.patch.signature import calc_truss_signature
+from truss.truss_handle.patch.signature import calc_truss_signature
 
 
 def test_calc_truss_signature(custom_model_truss_dir):

truss/tests/patch/test_truss_dir_patch_applier.py

@@ -2,17 +2,18 @@ import logging
 from pathlib import Path
 
 import yaml
-from truss.patch.truss_dir_patch_applier import TrussDirPatchApplier
-from truss.server.control.patch.types import (
+from truss.base.truss_config import TrussConfig
+from truss.templates.control.control.helpers.custom_types import (
     Action,
     ConfigPatch,
     ModelCodePatch,
+    PackagePatch,
     Patch,
     PatchType,
     PythonRequirementPatch,
     SystemPackagePatch,
 )
-from truss.truss_config import TrussConfig
+from truss.truss_handle.patch.truss_dir_patch_applier import TrussDirPatchApplier
 
 TEST_LOGGER = logging.getLogger("test_logger")
 
@@ -32,6 +33,23 @@ def test_model_code_patch(custom_model_truss_dir: Path):
     assert (custom_model_truss_dir / "model" / "model.py").read_text() == "test_content"
 
 
+def test_packages_patch(custom_model_truss_dir: Path):
+    applier = TrussDirPatchApplier(custom_model_truss_dir, TEST_LOGGER)
+    applier(
+        [
+            Patch(
+                type=PatchType.PACKAGE,
+                body=PackagePatch(
+                    action=Action.UPDATE, path="user_package.py", content="import sys"
+                ),
+            )
+        ]
+    )
+    assert (
+        custom_model_truss_dir / "packages" / "user_package.py"
+    ).read_text() == "import sys"
+
+
 def test_python_requirement_patch(custom_model_truss_dir: Path):
     req = "git+https://github.com/huggingface/transformers.git"
     applier = TrussDirPatchApplier(custom_model_truss_dir, TEST_LOGGER)
@@ -41,10 +59,7 @@ def test_python_requirement_patch(custom_model_truss_dir: Path):
         [
             Patch(
                 type=PatchType.PYTHON_REQUIREMENT,
-                body=PythonRequirementPatch(
-                    action=Action.ADD,
-                    requirement=req,
-                ),
+                body=PythonRequirementPatch(action=Action.ADD, requirement=req),
             ),
             Patch(
                 type=PatchType.CONFIG,
@@ -73,10 +88,7 @@ def test_system_requirement_patch(custom_model_truss_dir: Path):
             ),
             Patch(
                 type=PatchType.SYSTEM_PACKAGE,
-                body=SystemPackagePatch(
-                    action=Action.ADD,
-                    package="curl",
-                ),
+                body=SystemPackagePatch(action=Action.ADD, package="curl"),
             ),
         ]
     )

truss/tests/patch/test_types.py

@@ -1,5 +1,5 @@
-from truss.patch.signature import calc_truss_signature
-from truss.patch.types import TrussSignature
+from truss.truss_handle.patch.custom_types import TrussSignature
+from truss.truss_handle.patch.signature import calc_truss_signature
 
 
 def test_truss_signature_type(custom_model_truss_dir):

truss/tests/remote/baseten/test_api.py

@@ -4,9 +4,11 @@ import pytest
 import requests
 from requests import Response
 from truss.remote.baseten.api import BasetenApi
+from truss.remote.baseten.custom_types import ChainletDataAtomic, OracleData
 from truss.remote.baseten.error import ApiError
 
 
+@pytest.fixture
 def mock_auth_service():
     auth_service = mock.Mock()
     auth_token = mock.Mock(headers=lambda: {"Authorization": "Api-Key token"})
@@ -52,45 +54,62 @@ def mock_create_model_response():
     return response
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
-@mock.patch("requests.post", return_value=mock_successful_response())
-def test_post_graphql_query_success(mock_post, mock_auth_service):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
+def mock_create_development_model_response():
+    response = Response()
+    response.status_code = 200
+    response.json = mock.Mock(
+        return_value={"data": {"deploy_draft_truss": {"id": "12345"}}}
+    )
+    return response
+
+
+def mock_deploy_chain_deployment_response():
+    response = Response()
+    response.status_code = 200
+    response.json = mock.Mock(
+        return_value={
+            "data": {
+                "deploy_chain_atomic": {
+                    "chain_id": "12345",
+                    "chain_deployment_id": "54321",
+                    "entrypoint_model_id": "67890",
+                    "entrypoint_model_version_id": "09876",
+                }
+            }
+        }
+    )
+    return response
+
 
+@pytest.fixture
+def baseten_api(mock_auth_service):
+    return BasetenApi("https://app.test.com", mock_auth_service)
+
+
+@mock.patch("requests.post", return_value=mock_successful_response())
+def test_post_graphql_query_success(mock_post, baseten_api):
     response_data = {"data": {"status": "success"}}
 
-    result = api._post_graphql_query("sample_query_string")
+    result = baseten_api._post_graphql_query("sample_query_string")
 
     assert result == response_data
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_graphql_error_response())
-def test_post_graphql_query_error(mock_post, mock_auth_service):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
-
+def test_post_graphql_query_error(mock_post, baseten_api):
     with pytest.raises(ApiError):
-        api._post_graphql_query("sample_query_string")
+        baseten_api._post_graphql_query("sample_query_string")
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_unsuccessful_response())
-def test_post_requests_error(mock_post, mock_auth_service):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
+def test_post_requests_error(mock_post, baseten_api):
     with pytest.raises(requests.exceptions.HTTPError):
-        api._post_graphql_query("sample_query_string")
+        baseten_api._post_graphql_query("sample_query_string")
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_create_model_version_response())
-def test_create_model_version_from_truss(mock_post, mock_auth_service):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
-
-    api.create_model_version_from_truss(
+def test_create_model_version_from_truss(mock_post, baseten_api):
+    baseten_api.create_model_version_from_truss(
         "model_id",
         "s3key",
         "config_str",
@@ -98,8 +117,8 @@ def test_create_model_version_from_truss(mock_post, mock_auth_service):
         "client_version",
         False,
         False,
-        False,
         "deployment_name",
+        "production",
     )
 
     gql_mutation = mock_post.call_args[1]["data"]["query"]
@@ -109,27 +128,22 @@ def test_create_model_version_from_truss(mock_post, mock_auth_service):
     assert 'semver_bump: "semver_bump"' in gql_mutation
     assert 'client_version: "client_version"' in gql_mutation
     assert "is_trusted: false" in gql_mutation
-    assert "promote_after_deploy: false" in gql_mutation
     assert "scale_down_old_production: true" in gql_mutation
     assert 'name: "deployment_name"' in gql_mutation
+    assert 'environment_name: "production"' in gql_mutation
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_create_model_version_response())
 def test_create_model_version_from_truss_does_not_send_deployment_name_if_not_specified(
-    mock_post, mock_auth_service
+    mock_post, baseten_api
 ):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
-
-    api.create_model_version_from_truss(
+    baseten_api.create_model_version_from_truss(
         "model_id",
         "s3key",
         "config_str",
         "semver_bump",
         "client_version",
         True,
-        True,
         False,
         deployment_name=None,
     )
@@ -141,20 +155,16 @@ def test_create_model_version_from_truss_does_not_send_deployment_name_if_not_sp
     assert 'semver_bump: "semver_bump"' in gql_mutation
     assert 'client_version: "client_version"' in gql_mutation
     assert "is_trusted: true" in gql_mutation
-    assert "promote_after_deploy: true" in gql_mutation
     assert "scale_down_old_production: true" in gql_mutation
-    assert "name: " not in gql_mutation
+    assert " name: " not in gql_mutation
+    assert "environment_name: " not in gql_mutation
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_create_model_version_response())
 def test_create_model_version_from_truss_does_not_scale_old_prod_to_zero_if_keep_previous_prod_settings(
-    mock_post, mock_auth_service
+    mock_post, baseten_api
 ):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
-
-    api.create_model_version_from_truss(
+    baseten_api.create_model_version_from_truss(
         "model_id",
         "s3key",
         "config_str",
@@ -162,8 +172,8 @@ def test_create_model_version_from_truss_does_not_scale_old_prod_to_zero_if_keep
         "client_version",
         True,
         True,
-        True,
         deployment_name=None,
+        environment="staging",
     )
 
     gql_mutation = mock_post.call_args[1]["data"]["query"]
@@ -173,25 +183,21 @@ def test_create_model_version_from_truss_does_not_scale_old_prod_to_zero_if_keep
     assert 'semver_bump: "semver_bump"' in gql_mutation
     assert 'client_version: "client_version"' in gql_mutation
     assert "is_trusted: true" in gql_mutation
-    assert "promote_after_deploy: true" in gql_mutation
     assert "scale_down_old_production: false" in gql_mutation
-    assert "name: " not in gql_mutation
+    assert " name: " not in gql_mutation
+    assert 'environment_name: "staging"' in gql_mutation
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_create_model_response())
-def test_create_model_from_truss(mock_post, mock_auth_service):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
-
-    api.create_model_from_truss(
+def test_create_model_from_truss(mock_post, baseten_api):
+    baseten_api.create_model_from_truss(
         "model_name",
         "s3key",
         "config_str",
         "semver_bump",
         "client_version",
-        False,
-        "deployment_name",
+        is_trusted=False,
+        deployment_name="deployment_name",
     )
 
     gql_mutation = mock_post.call_args[1]["data"]["query"]
@@ -204,21 +210,17 @@ def test_create_model_from_truss(mock_post, mock_auth_service):
     assert 'version_name: "deployment_name"' in gql_mutation
 
 
-@mock.patch("truss.remote.baseten.auth.AuthService")
 @mock.patch("requests.post", return_value=mock_create_model_response())
 def test_create_model_from_truss_does_not_send_deployment_name_if_not_specified(
-    mock_post, mock_auth_service
+    mock_post, baseten_api
 ):
-    api_url = "https://test.com/api"
-    api = BasetenApi(api_url, mock_auth_service)
-
-    api.create_model_from_truss(
+    baseten_api.create_model_from_truss(
         "model_name",
         "s3key",
         "config_str",
         "semver_bump",
         "client_version",
-        True,
+        is_trusted=True,
         deployment_name=None,
     )
 
@@ -230,3 +232,96 @@ def test_create_model_from_truss_does_not_send_deployment_name_if_not_specified(
     assert 'client_version: "client_version"' in gql_mutation
     assert "is_trusted: true" in gql_mutation
     assert "version_name: " not in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_create_model_response())
+def test_create_model_from_truss_with_allow_truss_download(mock_post, baseten_api):
+    baseten_api.create_model_from_truss(
+        "model_name",
+        "s3key",
+        "config_str",
+        "semver_bump",
+        "client_version",
+        is_trusted=True,
+        allow_truss_download=False,
+    )
+
+    gql_mutation = mock_post.call_args[1]["data"]["query"]
+    assert 'name: "model_name"' in gql_mutation
+    assert 's3_key: "s3key"' in gql_mutation
+    assert 'config: "config_str"' in gql_mutation
+    assert 'semver_bump: "semver_bump"' in gql_mutation
+    assert 'client_version: "client_version"' in gql_mutation
+    assert "is_trusted: true" in gql_mutation
+    assert "allow_truss_download: false" in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_create_development_model_response())
+def test_create_development_model_from_truss_with_allow_truss_download(
+    mock_post, baseten_api
+):
+    baseten_api.create_development_model_from_truss(
+        "model_name",
+        "s3key",
+        "config_str",
+        "client_version",
+        is_trusted=True,
+        allow_truss_download=False,
+    )
+
+    gql_mutation = mock_post.call_args[1]["data"]["query"]
+    assert 'name: "model_name"' in gql_mutation
+    assert 's3_key: "s3key"' in gql_mutation
+    assert 'config: "config_str"' in gql_mutation
+    assert 'client_version: "client_version"' in gql_mutation
+    assert "is_trusted: true" in gql_mutation
+    assert "allow_truss_download: false" in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_deploy_chain_deployment_response())
+def test_deploy_chain_deployment(mock_post, baseten_api):
+    baseten_api.deploy_chain_atomic(
+        environment="production",
+        chain_id="chain_id",
+        dependencies=[],
+        entrypoint=ChainletDataAtomic(
+            name="chainlet-1",
+            oracle=OracleData(
+                model_name="model-1",
+                s3_key="s3-key-1",
+                encoded_config_str="encoded-config-str-1",
+                is_trusted=True,
+            ),
+        ),
+    )
+
+    gql_mutation = mock_post.call_args[1]["data"]["query"]
+
+    assert 'environment: "production"' in gql_mutation
+    assert 'chain_id: "chain_id"' in gql_mutation
+    assert "dependencies:" in gql_mutation
+    assert "entrypoint:" in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_deploy_chain_deployment_response())
+def test_deploy_chain_deployment_no_environment(mock_post, baseten_api):
+    baseten_api.deploy_chain_atomic(
+        chain_id="chain_id",
+        dependencies=[],
+        entrypoint=ChainletDataAtomic(
+            name="chainlet-1",
+            oracle=OracleData(
+                model_name="model-1",
+                s3_key="s3-key-1",
+                encoded_config_str="encoded-config-str-1",
+                is_trusted=True,
+            ),
+        ),
+    )
+
+    gql_mutation = mock_post.call_args[1]["data"]["query"]
+
+    assert 'chain_id: "chain_id"' in gql_mutation
+    assert "environment" not in gql_mutation
+    assert "dependencies:" in gql_mutation
+    assert "entrypoint:" in gql_mutation

truss/tests/remote/baseten/test_auth.py

@@ -9,7 +9,8 @@ def test_api_key():
     assert key.header() == {"Authorization": "Api-Key test_key"}
 
 
-def test_auth_service_no_key():
+def test_auth_service_no_key(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("BASETEN_API_KEY", raising=False)
     auth_service = AuthService()
     with pytest.raises(AuthorizationError):
         auth_service.authenticate()

truss/tests/remote/baseten/test_core.py

@@ -1,8 +1,13 @@
+import json
 from tempfile import NamedTemporaryFile
 from unittest.mock import MagicMock
 
+import pytest
+from truss.base.constants import PRODUCTION_ENVIRONMENT_NAME
+from truss.base.errors import ValidationError
 from truss.remote.baseten import core
 from truss.remote.baseten.api import BasetenApi
+from truss.remote.baseten.core import create_truss_service
 from truss.remote.baseten.error import ApiError
 
 
@@ -35,31 +40,23 @@ def test_upload_truss():
     core.multipart_upload_boto3 = MagicMock()
     core.multipart_upload_boto3.return_value = None
     test_file = NamedTemporaryFile()
-    assert core.upload_truss(api, test_file) == "key"
+    assert core.upload_truss(api, test_file, None) == "key"
 
 
 def test_get_dev_version_from_versions():
-    versions = [
-        {"id": "1", "is_draft": False},
-        {"id": "2", "is_draft": True},
-    ]
+    versions = [{"id": "1", "is_draft": False}, {"id": "2", "is_draft": True}]
     dev_version = core.get_dev_version_from_versions(versions)
     assert dev_version["id"] == "2"
 
 
 def test_get_dev_version_from_versions_error():
-    versions = [
-        {"id": "1", "is_draft": False},
-    ]
+    versions = [{"id": "1", "is_draft": False}]
     dev_version = core.get_dev_version_from_versions(versions)
     assert dev_version is None
 
 
 def test_get_dev_version():
-    versions = [
-        {"id": "1", "is_draft": False},
-        {"id": "2", "is_draft": True},
-    ]
+    versions = [{"id": "1", "is_draft": False}, {"id": "2", "is_draft": True}]
     api = MagicMock()
     api.get_model.return_value = {"model": {"versions": versions}}
 
@@ -84,3 +81,154 @@ def test_get_prod_version_from_versions_error():
     ]
     prod_version = core.get_prod_version_from_versions(versions)
     assert prod_version is None
+
+
+@pytest.mark.parametrize("environment", [None, PRODUCTION_ENVIRONMENT_NAME])
+def test_create_truss_service_handles_eligible_environment_values(environment):
+    api = MagicMock()
+    return_value = {"id": "id", "version_id": "model_version_id"}
+    api.create_model_from_truss.return_value = return_value
+    model_id, model_version_id = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        is_trusted=False,
+        preserve_previous_prod_deployment=False,
+        is_draft=False,
+        model_id=None,
+        deployment_name="deployment_name",
+        environment=environment,
+    )
+    assert model_id == return_value["id"]
+    assert model_version_id == return_value["version_id"]
+    api.create_model_from_truss.assert_called_once()
+
+
+@pytest.mark.parametrize("model_id", ["some_model_id", None])
+def test_create_truss_services_handles_is_draft(model_id):
+    api = MagicMock()
+    return_value = {"id": "id", "version_id": "model_version_id"}
+    api.create_development_model_from_truss.return_value = return_value
+    model_id, model_version_id = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        is_trusted=False,
+        preserve_previous_prod_deployment=False,
+        is_draft=True,
+        model_id=model_id,
+        deployment_name="deployment_name",
+    )
+    assert model_id == return_value["id"]
+    assert model_version_id == return_value["version_id"]
+    api.create_development_model_from_truss.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "inputs",
+    [
+        {
+            "environment": None,
+            "deployment_name": "some deployment",
+            "is_trusted": True,
+            "preserve_previous_prod_deployment": False,
+        },
+        {
+            "environment": PRODUCTION_ENVIRONMENT_NAME,
+            "deployment_name": None,
+            "is_trusted": True,
+            "preserve_previous_prod_deployment": False,
+        },
+        {
+            "environment": "staging",
+            "deployment_name": "some_deployment_name",
+            "is_trusted": False,
+            "preserve_previous_prod_deployment": True,
+        },
+    ],
+)
+def test_create_truss_service_handles_existing_model(inputs):
+    api = MagicMock()
+    return_value = {"id": "model_version_id"}
+    api.create_model_version_from_truss.return_value = return_value
+    model_id, model_version_id = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        is_draft=False,
+        model_id="model_id",
+        **inputs,
+    )
+
+    assert model_id == "model_id"
+    assert model_version_id == return_value["id"]
+    api.create_model_version_from_truss.assert_called_once()
+    _, kwargs = api.create_model_version_from_truss.call_args
+    for k, v in inputs.items():
+        assert kwargs[k] == v
+
+
+@pytest.mark.parametrize("allow_truss_download", [True, False])
+@pytest.mark.parametrize("is_draft", [True, False])
+def test_create_truss_service_handles_allow_truss_download_for_new_models(
+    is_draft, allow_truss_download
+):
+    api = MagicMock()
+    return_value = {"id": "id", "version_id": "model_version_id"}
+    api.create_model_from_truss.return_value = return_value
+    api.create_development_model_from_truss.return_value = return_value
+
+    model_id = None
+    model_id, model_version_id = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        is_trusted=False,
+        preserve_previous_prod_deployment=False,
+        is_draft=is_draft,
+        model_id=model_id,
+        deployment_name="deployment_name",
+        allow_truss_download=allow_truss_download,
+    )
+    assert model_id == return_value["id"]
+    assert model_version_id == return_value["version_id"]
+
+    create_model_mock = (
+        api.create_development_model_from_truss
+        if is_draft
+        else api.create_model_from_truss
+    )
+    create_model_mock.assert_called_once()
+    _, kwargs = create_model_mock.call_args
+    assert kwargs["allow_truss_download"] is allow_truss_download
+
+
+def test_validate_truss_config():
+    def mock_validate_truss(client_version, config):
+        if config == {}:
+            return {"success": True, "details": json.dumps({})}
+        elif "hi" in config:
+            return {"success": False, "details": json.dumps({"errors": ["error"]})}
+        else:
+            return {
+                "success": False,
+                "details": json.dumps({"errors": ["error", "and another one"]}),
+            }
+
+    api = MagicMock()
+    api.validate_truss.side_effect = mock_validate_truss
+
+    assert core.validate_truss_config(api, {}) is None
+    with pytest.raises(
+        ValidationError, match="Validation failed with the following errors:\n  error"
+    ):
+        core.validate_truss_config(api, {"hi": "hi"})
+    with pytest.raises(
+        ValidationError,
+        match="Validation failed with the following errors:\n  error\n  and another one",
+    ):
+        core.validate_truss_config(api, {"should_error": "hi"})