feat: add automation context implementation

tests/test_function.py

@@ -0,0 +1,173 @@
+"""Run integration tests with a speckle server."""
+import os
+import secrets
+import string
+from pathlib import Path
+
+import pytest
+from gql import gql
+from specklepy.api import operations
+from specklepy.api.client import SpeckleClient
+from specklepy.objects.base import Base
+from specklepy.transports.server import ServerTransport
+
+from automate_sdk import (
+    AutomationContext,
+    AutomationRunData,
+    AutomationStatus,
+    run_function,
+)
+from main import FunctionInputs, automate_function
+
+
+def crypto_random_string(length: int) -> str:
+    """Generate a semi crypto random string of a given length."""
+    alphabet = string.ascii_letters + string.digits
+    return "".join(secrets.choice(alphabet) for _ in range(length))
+
+
+def register_new_automation(
+    project_id: str,
+    model_id: str,
+    speckle_client: SpeckleClient,
+    automation_id: str,
+    automation_name: str,
+    automation_revision_id: str,
+):
+    """Register a new automation in the speckle server."""
+    query = gql(
+        """
+        mutation CreateAutomation(
+            $projectId: String! 
+            $modelId: String! 
+            $automationName: String!
+            $automationId: String! 
+            $automationRevisionId: String!
+        ) {
+                automationMutations {
+                    create(
+                        input: {
+                            projectId: $projectId
+                            modelId: $modelId
+                            automationName: $automationName 
+                            automationId: $automationId
+                            automationRevisionId: $automationRevisionId
+                        }
+                    )
+                }
+            }
+        """
+    )
+    params = {
+        "projectId": project_id,
+        "modelId": model_id,
+        "automationName": automation_name,
+        "automationId": automation_id,
+        "automationRevisionId": automation_revision_id,
+    }
+    speckle_client.httpclient.execute(query, params)
+
+
+@pytest.fixture()
+def speckle_token() -> str:
+    """Provide a speckle token for the test suite."""
+    env_var = "SPECKLE_TOKEN"
+    token = os.getenv(env_var)
+    if not token:
+        raise ValueError(f"Cannot run tests without a {env_var} environment variable")
+    return token
+
+
+@pytest.fixture()
+def speckle_server_url() -> str:
+    """Provide a speckle server url for the test suite, default to localhost."""
+    return os.getenv("SPECKLE_SERVER_URL", "http://127.0.0.1:3000")
+
+
+@pytest.fixture()
+def test_client(speckle_server_url: str, speckle_token: str) -> SpeckleClient:
+    """Initialize a SpeckleClient for testing."""
+    test_client = SpeckleClient(
+        speckle_server_url, speckle_server_url.startswith("https")
+    )
+    test_client.authenticate_with_token(speckle_token)
+    return test_client
+
+
+@pytest.fixture()
+def test_object() -> Base:
+    """Create a Base model for testing."""
+    root_object = Base()
+    root_object.foo = "bar"
+    return root_object
+
+
+@pytest.fixture()
+def automation_run_data(
+    test_object: Base, test_client: SpeckleClient, speckle_server_url: str
+) -> AutomationRunData:
+    """Set up an automation context for testing."""
+    project_id = test_client.stream.create("Automate function e2e test")
+    branch_name = "main"
+
+    model = test_client.branch.get(project_id, branch_name, commits_limit=1)
+    model_id: str = model.id
+
+    root_obj_id = operations.send(
+        test_object, [ServerTransport(project_id, test_client)]
+    )
+    version_id = test_client.commit.create(project_id, root_obj_id)
+
+    automation_name = crypto_random_string(10)
+    automation_id = crypto_random_string(10)
+    automation_revision_id = crypto_random_string(10)
+
+    register_new_automation(
+        project_id,
+        model_id,
+        test_client,
+        automation_id,
+        automation_name,
+        automation_revision_id,
+    )
+
+    automation_run_id = crypto_random_string(10)
+    function_id = crypto_random_string(10)
+    function_revision = crypto_random_string(10)
+    return AutomationRunData(
+        project_id=project_id,
+        model_id=model_id,
+        branch_name=branch_name,
+        version_id=version_id,
+        speckle_server_url=speckle_server_url,
+        automation_id=automation_id,
+        automation_revision_id=automation_revision_id,
+        automation_run_id=automation_run_id,
+        function_id=function_id,
+        function_revision=function_revision,
+    )
+
+
+def test_function_run(automation_run_data: AutomationRunData, speckle_token: str):
+    """Run an integration test for the automate function."""
+    automate_sdk = run_function(
+        automate_function,
+        automation_run_data,
+        speckle_token,
+        FunctionInputs(forbidden_speckle_type="Base"),
+    )
+
+    assert automate_sdk.run_status == AutomationStatus.FAILED
+
+
+def test_file_uploads(automation_run_data: AutomationRunData, speckle_token: str):
+    """Test file store capabilities of the automate sdk."""
+    automate_context = AutomationContext.initialize(automation_run_data, speckle_token)
+
+    path = Path(f"./{crypto_random_string(10)}").resolve()
+    path.write_text("foobar")
+
+    automate_context.store_file_result(path)
+
+    os.remove(path)
+    assert len(automate_context._automation_result.blobs) == 1