Merge pull request #318 from specklesystems/gergo/markLogs

feat: reduce log message context for object results

.circleci/config.yml

@@ -8,7 +8,7 @@ jobs:
   test:
     machine:
       image: ubuntu-2204:2023.02.1
-      docker_layer_caching: true
+      docker_layer_caching: false
     resource_class: medium
     parameters:
       tag:

docker-compose.yml

@@ -33,7 +33,7 @@ services:
       retries: 30
 
   minio:
-    image: "minio/minio"
+    image: "minio/minio:RELEASE.2023-10-25T06-33-25Z"
     command: server /data --console-address ":9001"
     restart: always
     volumes:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "specklepy"
-version = "2.9.1"
+version = "2.17.8"
 description = "The Python SDK for Speckle 2.0"
 readme = "README.md"
 authors = ["Speckle Systems <devops@speckle.systems>"]

src/speckle_automate/__init__.py

@@ -6,8 +6,8 @@ from speckle_automate.schema import (
     AutomationResult,
     AutomationRunData,
     AutomationStatus,
-    ObjectResult,
     ObjectResultLevel,
+    ResultCase,
 )
 
 __all__ = [
@@ -16,7 +16,7 @@ __all__ = [
     "AutomationStatus",
     "AutomationResult",
     "AutomationRunData",
-    "ObjectResult",
+    "ResultCase",
     "ObjectResultLevel",
     "run_function",
     "execute_automate_function",

src/speckle_automate/automation_context.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 import time
-from typing import Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import httpx
 from gql import gql
@@ -19,8 +19,8 @@ from speckle_automate.schema import (
     AutomationResult,
     AutomationRunData,
     AutomationStatus,
-    ObjectResult,
     ObjectResultLevel,
+    ResultCase,
 )
 
 
@@ -84,6 +84,11 @@ class AutomationContext:
         """Get the status of the automation run."""
         return self._automation_result.run_status
 
+    @property
+    def status_message(self) -> Optional[str]:
+        """Get the current status message."""
+        return self._automation_result.status_message
+
     def elapsed(self) -> float:
         """Return the elapsed time in seconds since the initialization time."""
         return time.perf_counter() - self._init_time
@@ -99,14 +104,14 @@ class AutomationContext:
             commit.referencedObject, self._server_transport, self._memory_transport
         )
         print(
-            f"It took {self.elapsed():2f} seconds to receive",
+            f"It took {self.elapsed():.2f} seconds to receive",
             f" the speckle version {self.automation_run_data.version_id}",
         )
         return base
 
     def create_new_version_in_project(
-        self, root_object: Base, model_id: str, version_message: str = ""
-    ) -> None:
+        self, root_object: Base, model_name: str, version_message: str = ""
+    ) -> Tuple[str, str]:
         """Save a base model to a new version on the project.
 
         Args:
@@ -114,15 +119,29 @@ class AutomationContext:
             model_id (str): For now please use a `branchName`!
             version_message (str): The message for the new version.
         """
-        if model_id == self.automation_run_data.model_id:
+
+        if model_name == self.automation_run_data.branch_name:
             raise ValueError(
-                f"The target model id: {model_id} cannot match the model id"
-                f" that triggered this automation: {self.automation_run_data.model_id}"
+                f"The target model: {model_name} cannot match the model"
+                f" that triggered this automation:"
+                f" {self.automation_run_data.model_id} /"
+                f" {self.automation_run_data.branch_name}"
             )
 
-        branch = self._get_model(model_id)
-        if not branch.name:
-            raise ValueError(f"The model {model_id} has no name.")
+        branch = self.speckle_client.branch.get(
+            self.automation_run_data.project_id, model_name, 1
+        )
+        # we just check if it exists
+        if (not branch) or isinstance(branch, SpeckleException):
+            branch_create = self.speckle_client.branch.create(
+                self.automation_run_data.project_id,
+                model_name,
+            )
+            if isinstance(branch_create, Exception):
+                raise branch_create
+            model_id = branch_create
+        else:
+            model_id = branch.id
 
         root_object_id = operations.send(
             root_object,
@@ -133,7 +152,7 @@ class AutomationContext:
         version_id = self.speckle_client.commit.create(
             stream_id=self.automation_run_data.project_id,
             object_id=root_object_id,
-            branch_name=branch.name,
+            branch_name=model_name,
             message=version_message,
             source_application="SpeckleAutomate",
         )
@@ -142,24 +161,35 @@ class AutomationContext:
             raise version_id
 
         self._automation_result.result_versions.append(version_id)
+        return model_id, version_id
 
-    def _get_model(self, model_id: str) -> Branch:
-        query = gql(
-            """
-            query ProjectModel($projectId: String!, $modelId: String!){
-                project(id: $projectId) {
-                    model(id: $modelId) {
-                        name
-                        id
-                        description
-                    }
-                }
-            }
-        """
+    @property
+    def context_view(self) -> Optional[str]:
+        return self._automation_result.result_view
+
+    def set_context_view(
+        self,
+        # f"{model_id}@{version_id} or {model_id} "
+        resource_ids: Optional[List[str]] = None,
+        include_source_model_version: bool = True,
+    ) -> None:
+        link_resources = (
+            [
+                f"{self.automation_run_data.model_id}@{self.automation_run_data.version_id}"
+            ]
+            if include_source_model_version
+            else []
+        )
+        if resource_ids:
+            link_resources.extend(resource_ids)
+        if not link_resources:
+            raise Exception(
+                "We do not have enough resource ids to compose a context view"
+            )
+        self._automation_result.result_view = (
+            f"/projects/{self.automation_run_data.project_id}"
+            f"/models/{','.join(link_resources)}"
         )
-        params = {"projectId": self.automation_run_data.project_id, "modelId": model_id}
-        response = self.speckle_client.httpclient.execute(query, params)
-        return Branch.model_validate(response["project"]["model"])
 
     def report_run_status(self) -> None:
         """Report the current run status to the project of this automation."""
@@ -171,8 +201,11 @@ class AutomationContext:
                 $automationRunId: String!,
                 $versionId: String!,
                 $functionId: String!,
+                $functionName: String!,
+                $functionLogo: String,
                 $runStatus: AutomationRunStatus!
                 $elapsed: Float!
+                $contextView: String
                 $resultVersionIds: [String!]!
                 $statusMessage: String
                 $objectResults: JSONObject
@@ -186,7 +219,10 @@ class AutomationContext:
                         functionRuns: [
                         {
                             functionId: $functionId
+                            functionName: $functionName
+                            functionLogo: $functionLogo
                             status: $runStatus,
+                            contextView: $contextView,
                             elapsed: $elapsed,
                             resultVersionIds: $resultVersionIds,
                             statusMessage: $statusMessage
@@ -201,22 +237,26 @@ class AutomationContext:
             object_results = {
                 "version": "1.0.0",
                 "values": {
-                    "speckleObjects": self._automation_result.model_dump(by_alias=True)[
+                    "objectResults": self._automation_result.model_dump(by_alias=True)[
                         "objectResults"
                     ],
-                    "blobs": self._automation_result.blobs,
+                    "blobIds": self._automation_result.blobs,
                 },
             }
         else:
             object_results = None
+
         params = {
             "automationId": self.automation_run_data.automation_id,
             "automationRevisionId": self.automation_run_data.automation_revision_id,
             "automationRunId": self.automation_run_data.automation_run_id,
             "versionId": self.automation_run_data.version_id,
             "functionId": self.automation_run_data.function_id,
+            "functionName": self.automation_run_data.function_name,
+            "functionLogo": self.automation_run_data.function_logo,
             "runStatus": self.run_status.value,
             "statusMessage": self._automation_result.status_message,
+            "contextView": self._automation_result.result_view,
             "elapsed": self.elapsed(),
             "resultVersionIds": self._automation_result.result_versions,
             "objectResults": object_results,
@@ -260,8 +300,9 @@ class AutomationContext:
         if len(upload_response.upload_results) != 1:
             raise ValueError("Expecting one upload result.")
 
-        for upload_result in upload_response.upload_results:
-            self._automation_result.blobs.append(upload_result.blob_id)
+        self._automation_result.blobs.extend(
+            [upload_result.blob_id for upload_result in upload_response.upload_results]
+        )
 
     def mark_run_failed(self, status_message: str) -> None:
         """Mark the current run a failure."""
@@ -279,28 +320,101 @@ class AutomationContext:
         self._automation_result.run_status = status
         self._automation_result.elapsed = duration
 
-        msg = f"Automation run {status.value} after {duration:2f} seconds."
+        msg = f"Automation run {status.value} after {duration:.2f} seconds."
         print("\n".join([msg, status_message]) if status_message else msg)
 
-    def add_object_error(self, object_id: str, error_cause: str) -> None:
-        """Add an error to a given Speckle object."""
-        self._add_object_result(object_id, ObjectResultLevel.ERROR, error_cause)
-
-    def add_object_warning(self, object_id: str, warning: str) -> None:
-        """Add a warning to a given Speckle object."""
-        self._add_object_result(object_id, ObjectResultLevel.WARNING, warning)
-
-    def add_object_info(self, object_id: str, info: str) -> None:
-        """Add an info message to a given Speckle object."""
-        self._add_object_result(object_id, ObjectResultLevel.INFO, info)
-
-    def _add_object_result(
-        self, object_id: str, level: ObjectResultLevel, status_message: str
+    def attach_error_to_objects(
+        self,
+        category: str,
+        object_ids: Union[str, List[str]],
+        message: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        visual_overrides: Optional[Dict[str, Any]] = None,
     ) -> None:
+        """Add a new error case to the run results.
+
+        If the error cause has already created an error case,
+        the error will be extended with a new case refering to the causing objects.
+        Args:
+            error_tag (str): A short tag for the error type.
+            causing_object_ids (str[]): A list of object_id-s that are causing the error
+            error_messagge (Optional[str]): Optional error message.
+            metadata: User provided metadata key value pairs
+            visual_overrides: Case specific 3D visual overrides.
+        """
+        self.attach_result_to_objects(
+            ObjectResultLevel.ERROR,
+            category,
+            object_ids,
+            message,
+            metadata,
+            visual_overrides,
+        )
+
+    def attach_warning_to_objects(
+        self,
+        category: str,
+        object_ids: Union[str, List[str]],
+        message: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        visual_overrides: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Add a new warning case to the run results."""
+        self.attach_result_to_objects(
+            ObjectResultLevel.WARNING,
+            category,
+            object_ids,
+            message,
+            metadata,
+            visual_overrides,
+        )
+
+    def attach_info_to_objects(
+        self,
+        category: str,
+        object_ids: Union[str, List[str]],
+        message: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        visual_overrides: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Add a new info case to the run results."""
+        self.attach_result_to_objects(
+            ObjectResultLevel.INFO,
+            category,
+            object_ids,
+            message,
+            metadata,
+            visual_overrides,
+        )
+
+    def attach_result_to_objects(
+        self,
+        level: ObjectResultLevel,
+        category: str,
+        object_ids: Union[str, List[str]],
+        message: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        visual_overrides: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        if isinstance(object_ids, list):
+            if len(object_ids) < 1:
+                raise ValueError(
+                    f"Need atleast one object_id to report a(n) {level.value.upper()}"
+                )
+            id_list = object_ids
+        else:
+            id_list = [object_ids]
         print(
-            f"Object {object_id} was marked with {level.value.upper()}",
-            f" cause: {status_message}",
+            f"Created new {level.value.upper()}"
+            f" category: {category} caused by: {message}"
         )
-        self._automation_result.object_results[object_id].append(
-            ObjectResult(level=level, status_message=status_message)
+        self._automation_result.object_results.append(
+            ResultCase(
+                category=category,
+                level=level,
+                object_ids=id_list,
+                message=message,
+                metadata=metadata,
+                visual_overrides=visual_overrides,
+            )
         )

src/speckle_automate/helpers.py

@@ -0,0 +1,54 @@
+"""Some useful helpers for working with automation data."""
+import secrets
+import string
+
+from specklepy.api.client import SpeckleClient
+from gql import gql
+
+
+def register_new_automation(
+    speckle_client: SpeckleClient,
+    project_id: str,
+    model_id: str,
+    automation_id: str,
+    automation_name: str,
+    automation_revision_id: str,
+) -> bool:
+    """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,
+    }
+    return speckle_client.httpclient.execute(query, params)
+
+
+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)).lower()

src/speckle_automate/runner.py

@@ -66,6 +66,9 @@ def execute_automate_function(
             raise ValueError("Cannot get speckle token from arguments or environment")
 
         speckle_token = speckle_token if speckle_token else args[3]
+        automation_context = AutomationContext.initialize(
+            automation_run_data, speckle_token
+        )
 
         inputs = (
             input_schema.model_validate_json(function_inputs)
@@ -75,16 +78,14 @@ def execute_automate_function(
 
         if inputs:
             automation_context = run_function(
+                automation_context,
                 automate_function,  # type: ignore
-                automation_run_data,
-                speckle_token,
                 inputs,
             )
         else:
             automation_context = run_function(
+                automation_context,
                 automate_function,  # type: ignore
-                automation_run_data,
-                speckle_token,
             )
 
         exit_code = (
@@ -98,9 +99,8 @@ def execute_automate_function(
 
 @overload
 def run_function(
+    automation_context: AutomationContext,
     automate_function: AutomateFunction[T],
-    automation_run_data: Union[AutomationRunData, str],
-    speckle_token: str,
     inputs: T,
 ) -> AutomationContext:
     ...
@@ -108,23 +108,18 @@ def run_function(
 
 @overload
 def run_function(
+    automation_context: AutomationContext,
     automate_function: AutomateFunctionWithoutInputs,
-    automation_run_data: Union[AutomationRunData, str],
-    speckle_token: str,
 ) -> AutomationContext:
     ...
 
 
 def run_function(
+    automation_context: AutomationContext,
     automate_function: Union[AutomateFunction[T], AutomateFunctionWithoutInputs],
-    automation_run_data: Union[AutomationRunData, str],
-    speckle_token: str,
     inputs: Optional[T] = None,
 ) -> AutomationContext:
     """Run the provided function with the automate sdk context."""
-    automation_context = AutomationContext.initialize(
-        automation_run_data, speckle_token
-    )
     automation_context.report_run_status()
 
     try:
@@ -151,5 +146,7 @@ def run_function(
             "Function error. Check the automation run logs for details."
         )
     finally:
+        if not automation_context.context_view:
+            automation_context.set_context_view()
         automation_context.report_run_status()
         return automation_context

src/speckle_automate/schema.py

@@ -1,7 +1,6 @@
 """"""
-from collections import defaultdict
 from enum import Enum
-from typing import Optional, List, Dict
+from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 from stringcase import camelcase
@@ -27,7 +26,8 @@ class AutomationRunData(BaseModel):
     automation_run_id: str
 
     function_id: str
-    function_release: str
+    function_name: str
+    function_logo: Optional[str]
 
     model_config = ConfigDict(
         alias_generator=camelcase, populate_by_name=True, protected_namespaces=()
@@ -51,11 +51,15 @@ class ObjectResultLevel(str, Enum):
     ERROR = "ERROR"
 
 
-class ObjectResult(AutomateBase):
-    """An object level result."""
+class ResultCase(AutomateBase):
+    """A result case."""
 
+    category: str
     level: ObjectResultLevel
-    status_message: str
+    object_ids: List[str]
+    message: Optional[str]
+    metadata: Optional[Dict[str, Any]]
+    visual_overrides: Optional[Dict[str, Any]]
 
 
 class AutomationResult(AutomateBase):
@@ -67,7 +71,4 @@ class AutomationResult(AutomateBase):
     blobs: List[str] = Field(default_factory=list)
     run_status: AutomationStatus = AutomationStatus.RUNNING
     status_message: Optional[str] = None
-
-    object_results: Dict[str, List[ObjectResult]] = Field(
-        default_factory=lambda: defaultdict(list)  # typing: ignore
-    )
+    object_results: list[ResultCase] = Field(default_factory=list)

src/specklepy/api/client.py

@@ -1,14 +1,8 @@
-import re
-from typing import Dict
-from warnings import warn
-
 from deprecated import deprecated
-from gql import Client
 from gql.transport.exceptions import TransportServerError
 from gql.transport.requests import RequestsHTTPTransport
 from gql.transport.websockets import WebsocketsTransport
 
-from specklepy.api import resources
 from specklepy.api.credentials import Account, get_account_from_token
 from specklepy.api.resources import (
     user,
@@ -131,7 +125,9 @@ class SpeckleClient(CoreSpeckleClient):
         Arguments:
             token {str} -- an api token
         """
-        metrics.track(metrics.SDK, self.account, {"name": "Client Authenticate_deprecated"})
+        metrics.track(
+            metrics.SDK, self.account, {"name": "Client Authenticate_deprecated"}
+        )
         return super().authenticate(token)
 
     def authenticate_with_token(self, token: str) -> None:
@@ -143,7 +139,9 @@ class SpeckleClient(CoreSpeckleClient):
         Arguments:
             token {str} -- an api token
         """
-        metrics.track(metrics.SDK, self.account, {"name": "Client Authenticate With Token"})
+        metrics.track(
+            metrics.SDK, self.account, {"name": "Client Authenticate With Token"}
+        )
         return super().authenticate_with_token(token)
 
     def authenticate_with_account(self, account: Account) -> None:
@@ -155,5 +153,7 @@ class SpeckleClient(CoreSpeckleClient):
             account {Account} -- the account object which can be found with
             `get_default_account` or `get_local_accounts`
         """
-        metrics.track(metrics.SDK, self.account, {"name": "Client Authenticate With Account"})
+        metrics.track(
+            metrics.SDK, self.account, {"name": "Client Authenticate With Account"}
+        )
         return super().authenticate_with_account(account)

src/specklepy/api/credentials.py

@@ -9,12 +9,15 @@ from specklepy.logging import metrics
 from specklepy.logging.exceptions import SpeckleException
 from specklepy.transports.sqlite import SQLiteTransport
 
-# following imports seem to be unnecessary, but they need to stay 
+# following imports seem to be unnecessary, but they need to stay
 # to not break the scripts using these functions as non-core
-from specklepy.core.api.credentials import (Account, UserInfo, 
-                                            StreamWrapper, # deprecated 
-                                            get_local_accounts as core_get_local_accounts, 
-                                            get_account_from_token as core_get_account_from_token)
+from specklepy.core.api.credentials import (
+    Account,
+    UserInfo,
+    StreamWrapper,  # deprecated
+    get_local_accounts as core_get_local_accounts,
+    get_account_from_token as core_get_account_from_token,
+)
 
 
 def get_local_accounts(base_path: Optional[str] = None) -> List[Account]:
@@ -35,11 +38,12 @@ def get_local_accounts(base_path: Optional[str] = None) -> List[Account]:
             (acc for acc in accounts if acc.isDefault),
             accounts[0] if accounts else None,
         ),
-        {"name": "Get Local Accounts"}
+        {"name": "Get Local Accounts"},
     )
 
     return accounts
 
+
 def get_default_account(base_path: Optional[str] = None) -> Optional[Account]:
     """
     Gets this environment's default account if any. If there is no default,
@@ -61,7 +65,8 @@ def get_default_account(base_path: Optional[str] = None) -> Optional[Account]:
     metrics.initialise_tracker(default)
 
     return default
-    
+
+
 def get_account_from_token(token: str, server_url: str = None) -> Account:
     """Gets the local account for the token if it exists
     Arguments:
@@ -73,5 +78,5 @@ def get_account_from_token(token: str, server_url: str = None) -> Account:
     """
     account = core_get_account_from_token(token, server_url)
 
-    metrics.track( metrics.SDK, account, {"name": "Get Account From Token"} )
+    metrics.track(metrics.SDK, account, {"name": "Get Account From Token"})
     return account

src/specklepy/core/api/wrapper.py

@@ -1,5 +1,6 @@
 from urllib.parse import unquote, urlparse
 from warnings import warn
+from gql import gql
 
 from specklepy.core.api.client import SpeckleClient
 from specklepy.core.api.credentials import (
@@ -81,13 +82,25 @@ class StreamWrapper:
                 " provided."
             )
 
+        # check for fe2 URL
+        if "/projects/" in parsed.path:
+            use_fe2 = True
+        else:
+            use_fe2 = False
+
         while segments:
             segment = segments.pop(0)
-            if segments and segment.lower() == "streams":
+            if segments and (
+                (use_fe2 is False and segment.lower() == "streams")
+                or (use_fe2 is True and segment.lower() == "projects")
+            ):
                 self.stream_id = segments.pop(0)
             elif segments and segment.lower() == "commits":
                 self.commit_id = segments.pop(0)
-            elif segments and segment.lower() == "branches":
+            elif segments and (
+                (use_fe2 is False and segment.lower() == "branches")
+                or (use_fe2 is True and segment.lower() == "models")
+            ):
                 self.branch_name = unquote(segments.pop(0))
             elif segments and segment.lower() == "objects":
                 self.object_id = segments.pop(0)
@@ -101,6 +114,38 @@ class StreamWrapper:
                     " provided."
                 )
 
+        if use_fe2 is True and self.branch_name is not None:
+            if "," in self.branch_name:
+                raise SpeckleException("Multi-model urls are not supported yet")
+
+            if "@" in self.branch_name:
+                model_id = self.branch_name.split("@")[0]
+                self.commit_id = self.branch_name.split("@")[1]
+            else:
+                model_id = self.branch_name
+
+            # get branch name
+            query = gql(
+                """
+                query Project($project_id: String!, $model_id: String!) {
+                    project(id: $project_id) {
+                        id
+                        model(id: $model_id) {
+                            name
+                        }
+                    }
+                }
+            """
+            )
+            self._client = self.get_client()
+            params = {"project_id": self.stream_id, "model_id": model_id}
+            project = self._client.httpclient.execute(query, params)
+
+            try:
+                self.branch_name = project["project"]["model"]["name"]
+            except KeyError as ke:
+                raise SpeckleException("Project model name is not found", ke)
+
         if not self.stream_id:
             raise SpeckleException(
                 f"Cannot parse {url} into a stream wrapper class - no stream id found."

src/specklepy/objects/other.py

@@ -1,7 +1,7 @@
 from typing import Any, List, Optional
 from deprecated import deprecated
 
-from specklepy.objects.geometry import Point, Vector
+from specklepy.objects.geometry import Point, Vector, Plane, Polyline
 
 from .base import Base
 
@@ -71,6 +71,18 @@ class DisplayStyle(Base, speckle_type=OTHER + "DisplayStyle"):
     lineweight: float = 0
 
 
+class Text(Base, speckle_type=OTHER + "Text"):
+    """
+    Text object to render it on viewer.
+    """
+    plane: Plane
+    value: str
+    height: float
+    rotation: float
+    displayValue: Optional[List[Polyline]] = None
+    richText: Optional[str] = None
+
+
 class Transform(
     Base,
     speckle_type=OTHER + "Transform",

tests/intergration/speckle_automate/test_automation_context.py

@@ -1,13 +1,12 @@
 """Run integration tests with a speckle server."""
 import os
-import secrets
-import string
 from pathlib import Path
 from typing import Dict
 
 import pytest
 from gql import gql
 from speckle_automate.schema import AutomateBase
+from speckle_automate.helpers import register_new_automation, crypto_random_string
 from specklepy.api import operations
 from specklepy.api.client import SpeckleClient
 from specklepy.objects.base import Base
@@ -21,67 +20,19 @@ from speckle_automate import (
 )
 
 
-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()
+@pytest.fixture
 def speckle_token(user_dict: Dict[str, str]) -> str:
     """Provide a speckle token for the test suite."""
     return user_dict["token"]
 
 
-@pytest.fixture()
+@pytest.fixture
 def speckle_server_url(host: str) -> str:
     """Provide a speckle server url for the test suite, default to localhost."""
     return f"http://{host}"
 
 
-@pytest.fixture()
+@pytest.fixture
 def test_client(speckle_server_url: str, speckle_token: str) -> SpeckleClient:
     """Initialize a SpeckleClient for testing."""
     test_client = SpeckleClient(speckle_server_url, use_ssl=False)
@@ -89,7 +40,7 @@ def test_client(speckle_server_url: str, speckle_token: str) -> SpeckleClient:
     return test_client
 
 
-@pytest.fixture()
+@pytest.fixture
 def test_object() -> Base:
     """Create a Base model for testing."""
     root_object = Base()
@@ -97,7 +48,7 @@ def test_object() -> Base:
     return root_object
 
 
-@pytest.fixture()
+@pytest.fixture
 def automation_run_data(
     test_object: Base, test_client: SpeckleClient, speckle_server_url: str
 ) -> AutomationRunData:
@@ -118,9 +69,9 @@ def automation_run_data(
     automation_revision_id = crypto_random_string(10)
 
     register_new_automation(
+        test_client,
         project_id,
         model_id,
-        test_client,
         automation_id,
         automation_name,
         automation_revision_id,
@@ -128,7 +79,7 @@ def automation_run_data(
 
     automation_run_id = crypto_random_string(10)
     function_id = crypto_random_string(10)
-    function_release = crypto_random_string(10)
+    function_name = f"automate test {crypto_random_string(3)}"
     return AutomationRunData(
         project_id=project_id,
         model_id=model_id,
@@ -139,10 +90,19 @@ def automation_run_data(
         automation_revision_id=automation_revision_id,
         automation_run_id=automation_run_id,
         function_id=function_id,
-        function_release=function_release,
+        function_name=function_name,
+        function_logo=None,
     )
 
 
+@pytest.fixture
+def automation_context(
+    automation_run_data: AutomationRunData, speckle_token: str
+) -> AutomationContext:
+    """Set up the run context."""
+    return AutomationContext.initialize(automation_run_data, speckle_token)
+
+
 def get_automation_status(
     project_id: str,
     model_id: str,
@@ -200,17 +160,18 @@ class FunctionInputs(AutomateBase):
 
 
 def automate_function(
-    automate_context: AutomationContext,
+    automation_context: AutomationContext,
     function_inputs: FunctionInputs,
 ) -> None:
     """Hey, trying the automate sdk experience here."""
-    version_root_object = automate_context.receive_version()
+    version_root_object = automation_context.receive_version()
 
     count = 0
     if version_root_object.speckle_type == function_inputs.forbidden_speckle_type:
         if not version_root_object.id:
             raise ValueError("Cannot operate on objects without their id's.")
-        automate_context.add_object_error(
+        automation_context.attach_error_to_objects(
+            "Forbidden speckle_type",
             version_root_object.id,
             "This project should not contain the type: "
             f"{function_inputs.forbidden_speckle_type}",
@@ -218,46 +179,107 @@ def automate_function(
         count += 1
 
     if count > 0:
-        automate_context.mark_run_failed(
+        automation_context.mark_run_failed(
             "Automation failed: "
             f"Found {count} object that have a forbidden speckle type: "
             f"{function_inputs.forbidden_speckle_type}"
         )
 
     else:
-        automate_context.mark_run_success("No forbidden types found.")
+        automation_context.mark_run_success("No forbidden types found.")
 
 
-def test_function_run(automation_run_data: AutomationRunData, speckle_token: str):
+def test_function_run(automation_context: AutomationContext) -> None:
     """Run an integration test for the automate function."""
     automation_context = run_function(
+        automation_context,
         automate_function,
-        automation_run_data,
-        speckle_token,
         FunctionInputs(forbidden_speckle_type="Base"),
     )
 
     assert automation_context.run_status == AutomationStatus.FAILED
     status = get_automation_status(
-        automation_run_data.project_id,
-        automation_run_data.model_id,
+        automation_context.automation_run_data.project_id,
+        automation_context.automation_run_data.model_id,
         automation_context.speckle_client,
     )
     assert status["status"] == automation_context.run_status
     status_message = status["automationRuns"][0]["functionRuns"][0]["statusMessage"]
-    assert status_message == automation_context._automation_result.status_message
+    assert status_message == automation_context.status_message
 
 
-def test_file_uploads(automation_run_data: AutomationRunData, speckle_token: str):
+@pytest.fixture
+def test_file_path():
+    path = Path(f"./{crypto_random_string(10)}").resolve()
+    yield path
+    os.remove(path)
+
+
+def test_file_uploads(
+    automation_run_data: AutomationRunData, speckle_token: str, test_file_path: Path
+):
     """Test file store capabilities of the automate sdk."""
     automation_context = AutomationContext.initialize(
         automation_run_data, speckle_token
     )
 
-    path = Path(f"./{crypto_random_string(10)}").resolve()
-    path.write_text("foobar")
+    test_file_path.write_text("foobar")
 
-    automation_context.store_file_result(path)
+    automation_context.store_file_result(test_file_path)
 
-    os.remove(path)
     assert len(automation_context._automation_result.blobs) == 1
+
+
+def test_create_version_in_project_raises_error_for_same_model(
+    automation_context: AutomationContext,
+) -> None:
+    with pytest.raises(ValueError):
+        automation_context.create_new_version_in_project(
+            Base(), automation_context.automation_run_data.branch_name
+        )
+
+
+def test_create_version_in_project(
+    automation_context: AutomationContext,
+) -> None:
+    root_object = Base()
+    root_object.foo = "bar"
+    model_id, version_id = automation_context.create_new_version_in_project(
+        root_object, "foobar"
+    )
+
+    assert model_id is not None
+    assert version_id is not None
+
+
+def test_set_context_view(automation_context: AutomationContext) -> None:
+    automation_context.set_context_view()
+
+    assert automation_context.context_view is not None
+    assert automation_context.context_view.endswith(
+        f"models/{automation_context.automation_run_data.model_id}@{automation_context.automation_run_data.version_id}"
+    )
+
+    automation_context.report_run_status()
+
+    automation_context._automation_result.result_view = None
+
+    dummy_context = "foo@bar"
+    automation_context.set_context_view([dummy_context])
+
+    assert automation_context.context_view is not None
+    assert automation_context.context_view.endswith(
+        f"models/{automation_context.automation_run_data.model_id}@{automation_context.automation_run_data.version_id},{dummy_context}"
+    )
+    automation_context.report_run_status()
+
+    automation_context._automation_result.result_view = None
+
+    dummy_context = "foo@baz"
+    automation_context.set_context_view(
+        [dummy_context], include_source_model_version=False
+    )
+
+    assert automation_context.context_view is not None
+    assert automation_context.context_view.endswith(f"models/{dummy_context}")
+    automation_context.report_run_status()

tests/intergration/test_wrapper.py

@@ -126,3 +126,24 @@ def test_wrapper_url_match(user_path) -> None:
     account = wrap.get_account()
 
     assert account.userInfo.email is None
+
+
+def test_parse_project():
+    wrap = StreamWrapper("https://latest.speckle.systems/projects/843d07eb10")
+    assert wrap.type == "stream"
+
+
+def test_parse_model():
+    wrap = StreamWrapper(
+        "https://latest.speckle.systems/projects/843d07eb10/models/d9eb4918c8"
+    )
+
+    assert wrap.branch_name == "building wrapper"
+    assert wrap.type == "branch"
+
+
+def test_parse_version():
+    wrap = StreamWrapper(
+        "https://latest.speckle.systems/projects/843d07eb10/models/4e7345c838@c42d5cbac1"
+    )
+    assert wrap.type == "commit"