truss 0.11.6rc102__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/tests/remote/baseten/test_remote.py

@@ -1,3 +1,5 @@
+from unittest import mock
+
 import pydantic
 import pytest
 import requests_mock
@@ -11,16 +13,17 @@ from truss.remote.baseten.core import (
 )
 from truss.remote.baseten.custom_types import ChainletDataAtomic, OracleData
 from truss.remote.baseten.error import RemoteError
-from truss.remote.baseten.remote import BasetenRemote
 from truss.truss_handle.truss_handle import TrussHandle
 
 _TEST_REMOTE_URL = "http://test_remote.com"
 _TEST_REMOTE_GRAPHQL_PATH = "http://test_remote.com/graphql/"
 
-
-@pytest.fixture
-def remote():
-    return BasetenRemote(_TEST_REMOTE_URL, "api_key")
+TRUSS_RC_CONTENT = """
+[baseten]
+remote_provider = baseten
+api_key = test_key
+remote_url = http://test.com
+""".strip()
 
 
 def assert_request_matches_expected_query(request, expected_query) -> None:
@@ -269,30 +272,41 @@ def test_push_raised_value_error_when_keep_previous_prod_settings_and_not_promot
             )
 
 
-def test_create_chain_with_no_publish(remote):
-    with requests_mock.Mocker() as m:
-        m.post(
-            _TEST_REMOTE_GRAPHQL_PATH,
-            [
-                {"json": {"data": {"chains": []}}},
-                {
-                    "json": {
-                        "data": {
-                            "deploy_chain_atomic": {
-                                "chain_deployment": {
-                                    "id": "new-chain-deployment-id",
-                                    "chain": {
-                                        "id": "new-chain-id",
-                                        "hostname": "hostname",
-                                    },
-                                }
-                            }
-                        }
-                    }
-                },
-            ],
+@pytest.mark.parametrize("deploy_timeout_minutes", [9, 1441])
+def test_push_raised_value_error_when_deploy_timeout_minutes_is_invalid(
+    deploy_timeout_minutes, custom_model_truss_dir_with_pre_and_post, remote
+):
+    th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
+
+    with pytest.raises(
+        ValueError,
+        match="deploy-timeout-minutes must be between 10 minutes and 1440 minutes \(24 hours\)",
+    ):
+        remote.push(
+            th,
+            "model_name",
+            th.truss_dir,
+            publish=True,
+            promote=False,
+            preserve_previous_prod_deployment=False,
+            deployment_name="dep_name",
+            deploy_timeout_minutes=deploy_timeout_minutes,
         )
 
+
+def test_create_chain_with_no_publish(remote):
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
+
+    with mock.patch.object(
+        remote.api, "get_chains", return_value=[]
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
         deployment_handle = create_chain_atomic(
             api=remote.api,
             chain_name="draft_chain",
@@ -310,64 +324,61 @@ def test_create_chain_with_no_publish(remote):
             environment=None,
         )
 
-        get_chains_graphql_request = m.request_history[0]
-        create_chain_graphql_request = m.request_history[1]
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
 
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_name"] == "draft_chain"
+        assert call_kwargs.get("is_draft") is True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
 
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
+        assert deployment_handle.chain_id == "new-chain-id"
+        assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
 
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        # Note that if publish=False and promote=True, we set publish to True and create
-        # a non-draft deployment
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_name: "draft_chain"
-                    is_draft: true
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
+
+def test_create_chain_no_existing_chain(remote):
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
+
+    with mock.patch.object(
+        remote.api, "get_chains", return_value=[]
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
+        deployment_handle = create_chain_atomic(
+            api=remote.api,
+            chain_name="new_chain",
+            entrypoint=ChainletDataAtomic(
+                name="chainlet-1",
+                oracle=OracleData(
+                    model_name="model-1",
+                    s3_key="s3-key-1",
+                    encoded_config_str="encoded-config-str-1",
+                ),
+            ),
+            dependencies=[],
+            truss_user_env=b10_types.TrussUserEnv.collect(),
+            is_draft=False,
+            environment=None,
         )
+
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
+
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_name"] == "new_chain"
+        assert call_kwargs.get("is_draft") is not True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
+
         assert deployment_handle.chain_id == "new-chain-id"
         assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
 
 
-def test_create_chain_no_existing_chain(remote):
+def test_create_chain_with_deployment_name(remote):
     with requests_mock.Mocker() as m:
         m.post(
             _TEST_REMOTE_GRAPHQL_PATH,
@@ -391,7 +402,8 @@ def test_create_chain_no_existing_chain(remote):
             ],
         )
 
-        deployment_handle = create_chain_atomic(
+        deployment_name = "chain-deployment"
+        create_chain_atomic(
             api=remote.api,
             chain_name="new_chain",
             entrypoint=ChainletDataAtomic(
@@ -406,94 +418,32 @@ def test_create_chain_no_existing_chain(remote):
             truss_user_env=b10_types.TrussUserEnv.collect(),
             is_draft=False,
             environment=None,
+            deployment_name=deployment_name,
         )
 
-        get_chains_graphql_request = m.request_history[0]
         create_chain_graphql_request = m.request_history[1]
 
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
-
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
-
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_name: "new_chain"
-                    is_draft: false
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
+        assert (
+            'deployment_name: "chain-deployment"'
+            in create_chain_graphql_request.json()["query"]
         )
 
-        assert deployment_handle.chain_id == "new-chain-id"
-        assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
-
 
 def test_create_chain_with_existing_chain_promote_to_environment_publish_false(remote):
-    with requests_mock.Mocker() as m:
-        m.post(
-            _TEST_REMOTE_GRAPHQL_PATH,
-            [
-                {
-                    "json": {
-                        "data": {
-                            "chains": [{"id": "old-chain-id", "name": "old_chain"}]
-                        }
-                    }
-                },
-                {
-                    "json": {
-                        "data": {
-                            "deploy_chain_atomic": {
-                                "chain_deployment": {
-                                    "id": "new-chain-deployment-id",
-                                    "chain": {
-                                        "id": "new-chain-id",
-                                        "hostname": "hostname",
-                                    },
-                                }
-                            }
-                        }
-                    }
-                },
-            ],
-        )
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
 
+    with mock.patch.object(
+        remote.api,
+        "get_chains",
+        return_value=[{"id": "old-chain-id", "name": "old_chain"}],
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
         deployment_handle = create_chain_atomic(
             api=remote.api,
             chain_name="old_chain",
@@ -511,95 +461,34 @@ def test_create_chain_with_existing_chain_promote_to_environment_publish_false(r
             environment="production",
         )
 
-        get_chains_graphql_request = m.request_history[0]
-        create_chain_graphql_request = m.request_history[1]
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
 
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
-
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
-
-        # Note that if publish=False and environment!=None, we set publish to True and create
-        # a non-draft deployment
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_id: "old-chain-id"
-                    environment: "production"
-                    is_draft: false
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
-        )
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_id"] == "old-chain-id"
+        assert call_kwargs["environment"] == "production"
+        assert call_kwargs.get("is_draft") is not True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
 
         assert deployment_handle.chain_id == "new-chain-id"
         assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
 
 
 def test_create_chain_existing_chain_publish_true_no_promotion(remote):
-    with requests_mock.Mocker() as m:
-        m.post(
-            _TEST_REMOTE_GRAPHQL_PATH,
-            [
-                {
-                    "json": {
-                        "data": {
-                            "chains": [{"id": "old-chain-id", "name": "old_chain"}]
-                        }
-                    }
-                },
-                {
-                    "json": {
-                        "data": {
-                            "deploy_chain_atomic": {
-                                "chain_deployment": {
-                                    "id": "new-chain-deployment-id",
-                                    "chain": {
-                                        "id": "new-chain-id",
-                                        "hostname": "hostname",
-                                    },
-                                }
-                            }
-                        }
-                    }
-                },
-            ],
-        )
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
 
+    with mock.patch.object(
+        remote.api,
+        "get_chains",
+        return_value=[{"id": "old-chain-id", "name": "old_chain"}],
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
         deployment_handle = create_chain_atomic(
             api=remote.api,
             chain_name="old_chain",
@@ -617,57 +506,13 @@ def test_create_chain_existing_chain_publish_true_no_promotion(remote):
             environment=None,
         )
 
-        get_chains_graphql_request = m.request_history[0]
-        create_chain_graphql_request = m.request_history[1]
-
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
-
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
 
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_id: "old-chain-id"
-                    is_draft: false
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
-        )
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_id"] == "old-chain-id"
+        assert call_kwargs.get("is_draft") is not True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
 
         assert deployment_handle.chain_id == "new-chain-id"
         assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
@@ -726,3 +571,86 @@ def test_push_raised_validation_error_for_extra_fields(tmp_path, remote):
             match="Extra fields not allowed: \[extra_field, who_am_i\]",
         ):
             remote.push(th, "model_name", th.truss_dir)
+
+
+def test_push_passes_deploy_timeout_minutes_to_create_truss_service(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+    mock_truss_handle,
+):
+    remote.push(
+        mock_truss_handle,
+        "model_name",
+        mock_truss_handle.truss_dir,
+        publish=True,
+        deploy_timeout_minutes=450,
+    )
+
+    mock_create_truss_service.assert_called_once()
+    _, kwargs = mock_create_truss_service.call_args
+    assert kwargs["deploy_timeout_minutes"] == 450
+
+
+def test_push_passes_none_deploy_timeout_minutes_when_not_specified(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+    mock_truss_handle,
+):
+    remote.push(
+        mock_truss_handle, "model_name", mock_truss_handle.truss_dir, publish=True
+    )
+
+    mock_create_truss_service.assert_called_once()
+    _, kwargs = mock_create_truss_service.call_args
+    assert kwargs.get("deploy_timeout_minutes") is None
+
+
+def test_push_integration_deploy_timeout_minutes_propagated(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+    mock_truss_handle,
+):
+    remote.push(
+        mock_truss_handle,
+        "model_name",
+        mock_truss_handle.truss_dir,
+        publish=True,
+        environment="staging",
+        deploy_timeout_minutes=750,
+    )
+
+    mock_create_truss_service.assert_called_once()
+    _, kwargs = mock_create_truss_service.call_args
+    assert kwargs["deploy_timeout_minutes"] == 750
+    assert kwargs["environment"] == "staging"
+
+
+def test_api_push_integration_deploy_timeout_minutes_propagated(
+    custom_model_truss_dir_with_pre_and_post,
+    mock_remote_factory,
+    temp_trussrc_dir,
+    mock_available_config_names,
+    mock_truss_handle,
+):
+    from truss.api import push
+
+    push(
+        str(mock_truss_handle.truss_dir),
+        remote="baseten",
+        model_name="test_model",
+        deploy_timeout_minutes=1200,
+    )
+
+    # Verify the remote.push was called with deploy_timeout_minutes
+    mock_remote_factory.push.assert_called_once()
+    _, push_kwargs = mock_remote_factory.push.call_args
+    assert push_kwargs.get("deploy_timeout_minutes") == 1200

truss/tests/remote/baseten/test_service.py

@@ -1,4 +1,7 @@
+from unittest.mock import MagicMock
+
 from truss.remote.baseten import service
+from truss.remote.baseten.core import ModelVersionHandle
 
 
 def test_model_invoke_url_prod():
@@ -65,3 +68,56 @@ def test_chain_logs_url():
         "https://app.baseten.co", "abc", "666", "543"
     )
     assert url == "https://app.baseten.co/chains/abc/logs/666/543"
+
+
+def test_predict_response_to_json():
+    """Test that predict method returns JSON response for normal dict result."""
+    # Create a mock BasetenService
+    mock_handle = MagicMock(spec=ModelVersionHandle)
+    mock_handle.model_id = "test-model"
+    mock_handle.version_id = "test-version"
+    mock_handle.hostname = "https://model-test.api.baseten.co"
+
+    mock_api = MagicMock()
+    mock_api.app_url = "https://app.baseten.co"
+
+    service_instance = service.BasetenService(
+        model_version_handle=mock_handle,
+        is_draft=False,
+        api_key="test-key",
+        service_url="https://test.com",
+        api=mock_api,
+    )
+
+    # Mock the _send_request method to return a successful JSON response
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"result": "success"}
+    service_instance._send_request = MagicMock(return_value=mock_response)
+
+    # Test predict method
+    result = service_instance.predict({"input": "test"})
+
+    # Verify that the JSON response is returned directly
+    assert result == {"result": "success"}
+
+    # Test non-dict response types below
+
+    # With integer response
+    mock_response.json.return_value = 42
+    result = service_instance.predict({"input": "test"})
+    assert result == 42
+
+    # With string response
+    mock_response.json.return_value = "success"
+    result = service_instance.predict({"input": "test"})
+    assert result == "success"
+
+    # With list response
+    mock_response.json.return_value = [1, 2, 3, 4]
+    result = service_instance.predict({"input": "test"})
+    assert result == [1, 2, 3, 4]
+
+    # With boolean response
+    mock_response.json.return_value = True
+    result = service_instance.predict({"input": "test"})
+    assert result is True

truss/tests/templates/control/control/conftest.py

@@ -0,0 +1,20 @@
+import sys
+from pathlib import Path
+
+
+# NB(nikhil): Slightly hacky helpers needed to set up the path so relative imports work as they do in real environments
+def setup_control_imports():
+    base_path = Path(__file__).parent.parent.parent.parent.parent
+    paths = [
+        base_path / "templates" / "control" / "control",
+        base_path / "templates",
+        base_path / "templates" / "shared",
+    ]
+
+    for path in paths:
+        if not path.exists():
+            raise FileNotFoundError(f"Expected control path does not exist: {path}")
+
+        path_str = str(path)
+        if path_str not in sys.path:
+            sys.path.insert(0, path_str)

truss/tests/templates/control/control/test_endpoints.py

@@ -8,6 +8,10 @@ from httpx_ws import AsyncWebSocketSession
 from httpx_ws import _exceptions as httpx_ws_exceptions
 from wsproto.events import BytesMessage, TextMessage
 
+from truss.tests.templates.control.control.conftest import setup_control_imports
+
+setup_control_imports()
+
 from truss.templates.control.control.endpoints import proxy_ws