Code Completion

Rules/actions.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Dict, List
+
+from speckle_automate import AutomationContext
+
+
+# Base class for defining actions to be taken on parameters in Speckle data.
+class ParameterAction(ABC):
+    """
+    A base class for creating actions that can be applied to parameters in
+    Speckle objects. This abstract class outlines the structure and mandates
+    the implementation of specific methods in derived classes.
+    """
+
+    def __init__(self) -> None:
+        # Dictionary for tracking affected parameters. Key: parent object's ID,
+        # Value: list of affected parameter names.
+        self.affected_parameters: Dict[str, List[str]] = defaultdict(list)
+
+    @abstractmethod
+    def apply(self, parameter: Dict[str, str], parent_object: Dict[str, str]) -> None:
+        """
+        Applies the specific logic of the action to a parameter.
+
+        Args:
+            parameter: The parameter to which the action is applied.
+            parent_object: The object that holds the parameter.
+        """
+        pass
+
+    @abstractmethod
+    def report(self, automate_context: AutomationContext) -> None:
+        """
+        Generates a report based on the results of applying the action.
+
+        Args:
+            automate_context: The context in which the automation is executed,
+            providing mechanisms for attaching results to the Speckle model.
+        """
+        pass
+
+# Further specific action classes can be defined here, inheriting from
+# ParameterAction and implementing the abstract methods 'apply' and 'report'.

Rules/checks.py

@@ -0,0 +1,136 @@
+# Required imports
+from typing import Callable, Dict, Union
+
+from specklepy.objects import Base
+
+
+# We're going to define a set of rules that will allow us to filter and
+# process parameters in our Speckle objects. These rules will be encapsulated
+# in a class called `ParameterRules`.
+
+
+class BaseObjectRules:
+    """
+    A collection of rules for processing parameters in Speckle objects.
+
+    This class provides static methods that return lambda functions. These
+    lambda functions serve as filters or conditions we can use in our main
+    processing logic. By encapsulating these rules, we can easily extend
+    or modify them in the future.
+    """
+
+    @staticmethod
+    def speckle_type_rule(desired_type: str) -> Callable[[Base], bool]:
+        """
+        Rule: Check if a parameter's speckle_type matches the desired type.
+        """
+        return (
+            lambda parameter: getattr(parameter, "speckle_type", None) == desired_type
+        )
+
+    @staticmethod
+    def forbidden_prefix_rule(given_prefix: str) -> Callable[[Base], bool]:
+        """
+        Rule: check if a parameter's name starts with a given prefix.
+
+        This is a simple check, but there could be more complex naming rules for parameters of
+        different types. For example, a rule that checks if a parameter's name starts with a given string
+        exists particularly within IFC where parameters are often prefixed with "Ifc" or "Pset".
+        """
+        return lambda parameter: parameter.name.startswith(given_prefix)
+
+    # This example Automate function is for prefixed parameter removal. Additional example rules below follow the same
+    # pattern, but with different logic. In some instances there is a strong coupling between the action and the checking
+    # logic, and in others there is a looser coupling. Which is why I have defined the actions separately from the
+    # checking logic.
+
+    @staticmethod
+    def has_missing_value(parameter: Union[Base, Dict[str, str]]) -> bool:
+        """
+        Rule: Missing Value Check.
+
+        The AEC industry often requires all parameters to have meaningful values.
+        This rule checks if a parameter is missing its value, potentially indicating
+        an oversight during data entry or transfer.
+        """
+        return not getattr(parameter, "value")
+
+    @staticmethod
+    def has_default_value(parameter: Dict[str, str]) -> bool:
+        """
+        Rule: Default Value Check.
+
+        Default values can sometimes creep into final datasets due to software defaults.
+        This rule identifies parameters that still have their default values, helping
+        to highlight areas where real, meaningful values need to be provided.
+        """
+        return parameter.get("value") == "Default"
+
+    @staticmethod
+    def parameter_exists(parameter_name: str, parent_object: Dict[str, str]) -> bool:
+        """
+        Rule: Parameter Existence Check.
+
+        For certain critical parameters, their mere presence (or lack thereof) is vital.
+        This rule verifies if a specific parameter exists within an object, allowing
+        teams to ensure that key data points are always present.
+        """
+        return parameter_name in parent_object.get("parameters", {})
+
+    @staticmethod
+    def is_category(category: str) -> Callable[[Base], bool]:
+        """
+        Rule: Category Check.
+
+        This rule checks if a parameter's category matches the desired category.
+        """
+        return lambda parameter: getattr(parameter, "category", None) == category
+
+    @staticmethod
+    def parameter_name_is(parameter_name: str) -> Callable[[Base], bool]:
+        """
+        Rule: Parameter Name Check.
+
+        This rule checks if a parameter's name matches the desired name.
+        """
+        return (
+            lambda parameter: getattr(parameter, "name") is not None
+            and parameter.name == parameter_name
+        )
+
+    @staticmethod
+    def parameter_value_startswith(prefix: str) -> Callable[[Base], bool]:
+        """
+        Rule: Parameter Name Starts With.
+
+        This rule checks if a parameter's name starts with a given prefix.
+        """
+        return lambda parameter: parameter.name.startswith(prefix)
+
+    @staticmethod
+    def is_revit_parameter(parameter: Union[Base, Dict[str, str]]):
+        """
+        Checks if a parameter is a Revit parameter.
+
+        This function checks if a parameter is a Revit parameter by checking if it
+        has a 'category' property.
+        """
+        return (
+            getattr(parameter, "speckle_type", None)
+            == "Objects.BuiltElements.Revit.Parameter"
+        )
+
+    @staticmethod
+    def evaluate_parameter(parameter, function_inputs):
+        """Evaluates a parameter and returns its evaluation state."""
+        if not BaseObjectRules.is_revit_parameter(parameter):
+            return None
+
+        if BaseObjectRules.has_missing_value(parameter):
+            return "missing"
+
+        value = getattr(parameter, "value", None)
+        if value is not None and value.startswith(function_inputs.single_rule):
+            return "valid"
+        else:
+            return "invalid"

Rules/traversal.py

@@ -0,0 +1,49 @@
+from specklepy.objects.graph_traversal.traversal import TraversalRule, GraphTraversal
+
+
+def get_data_traversal_rules() -> GraphTraversal:
+    """
+    Generates traversal rules for navigating Speckle data structures.
+
+    This function defines and returns traversal rules tailored for Speckle data.
+    These rules are used to navigate and extract specific data properties
+    within complex Speckle data hierarchies.
+
+    It defines two main rules:
+
+    1. `display_value_rule`:
+        - Targets objects that have properties named either "displayValue" or
+          "@displayValue".
+        - Specifically focuses on objects with a 'speckle_type' containing
+          "Geometry".
+        - For such objects, the function looks to traverse their 'elements'
+          or '@elements' properties.
+
+    2. `default_rule`:
+        - A more general rule that applies to all objects.
+        - It aims to traverse all member names of an object while avoiding
+          deprecated members (a potential enhancement for the future).
+
+    Returns:
+        GraphTraversal: A GraphTraversal instance initialized with the
+        defined rules.
+    """
+    display_value_property_aliases = {"displayValue", "@displayValue"}
+    elements_property_aliases = {"elements", "@elements"}
+
+    display_value_rule = TraversalRule(
+        [
+            lambda o: any(
+                getattr(o, alias, None) for alias in display_value_property_aliases
+            ),
+            lambda o: "Geometry" in o.speckle_type,
+        ],
+        lambda o: elements_property_aliases,
+    )
+
+    default_rule = TraversalRule(
+        [lambda _: True],
+        lambda o: o.get_member_names(),
+    )
+
+    return GraphTraversal([display_value_rule, default_rule])

Utilities/flatten.py

@@ -1,13 +1,90 @@
-"""Helper module for a simple speckle object tree flattening."""
+"""
+This helper module provides functions for flattening Speckle object trees and
+extracting base objects along with their transformations. It's designed for AEC
+professionals working with complex Speckle data structures.
+"""
 
 from collections.abc import Iterable
+from typing import Optional, Tuple, List
 
 from specklepy.objects import Base
+from specklepy.objects.other import Instance, Transform
 
 
-def flatten_base(base: Base) -> Iterable[Base]:
-    """Take a base and flatten it to an iterable of bases."""
-    if hasattr(base, "elements"):
-        for element in base["elements"]:
-            yield from flatten_base(element)
-    yield base
+def flatten_base(base: Base, parent_type: str = None) -> Iterable[Base]:
+    """
+    Flattens a Speckle object tree into an iterable of base objects.
+
+    Args:
+        base: The base object to flatten.
+        parent_type: The type of the parent object, if applicable.
+
+    Yields:
+        Base: A flattened base object, making complex hierarchies linear.
+    """
+    if isinstance(base, Base):
+        base["parent_type"] = parent_type
+
+    # Handle collections of elements in the base object
+    if hasattr(base, "elements") and base.elements:
+        try:
+            for element in base.elements:
+                yield from flatten_base(element, base.speckle_type)
+        except KeyError:
+            pass
+    # Handle older Revit-specific patterns with '@Lines'
+    elif hasattr(base, "@Lines"):
+        categories = base.get_dynamic_member_names()
+        for category in categories:
+            if category.startswith("@"):
+                category_object: Base = getattr(base, category)[0]
+                yield from flatten_base(category_object, category_object.speckle_type)
+    else:
+        yield base
+
+
+def extract_base_and_transform(
+    base: Base,
+    inherited_instance_id: Optional[str] = None,
+    transform_list: Optional[List[Transform]] = None,
+) -> Tuple[Base, str, Optional[List[Transform]]]:
+    """
+    Extracts `Base` objects and their transformations from Speckle data.
+
+    Args:
+        base: The starting point `Base` object for traversal.
+        inherited_instance_id: Inherited ID for objects without a unique one.
+        transform_list: List of transformations from parent to child objects.
+
+    Yields:
+        tuple: A `Base` object, its identifier, and applicable transforms.
+    """
+    current_id = getattr(base, "id", inherited_instance_id)
+    transform_list = transform_list or []
+
+    if isinstance(base, Instance):
+        if base.transform:
+            transform_list.append(base.transform)
+        if base.definition:
+            yield from extract_base_and_transform(
+                base.definition, current_id, transform_list.copy()
+            )
+    else:
+        yield base, current_id, transform_list
+
+        # Process 'elements' and '@elements' in the base object
+        elements_attr = getattr(base, "elements", []) or getattr(base, "@elements", [])
+        for element in elements_attr:
+            if isinstance(element, Base):
+                yield from extract_base_and_transform(
+                    element, current_id, transform_list.copy()
+                )
+
+        # Process '@'-prefixed properties in older Speckle data models
+        for attr_name in dir(base):
+            if attr_name.startswith("@"):
+                attr_value = getattr(base, attr_name)
+                if isinstance(attr_value, Base) and hasattr(attr_value, "elements"):
+                    yield from extract_base_and_transform(
+                        attr_value, current_id, transform_list.copy()
+                    )

Utilities/report.py

@@ -0,0 +1,161 @@
+import json
+from typing import List, Dict
+
+from fpdf import FPDF  # To install: `pip install fpdf2`
+
+
+def save_html_report(data: str, filename: str) -> None:
+    """
+    Saves HTML content as a file, handy for viewing in web browsers.
+
+    Args:
+        data (str): HTML content.
+        filename (str): File path to save HTML content.
+    """
+    with open(filename, "w") as file:
+        file.write(data)
+
+
+def save_json_report(
+    data: Dict[str, List[Dict[str, str]]],
+    filename: str,
+    single_category: str,
+    single_property: str,
+    single_value: str,
+) -> None:
+    """
+    Saves data as JSON. Ideal for data exchange or further processing.
+
+    Args:
+        data: The structured data to save.
+        filename: Where to save the JSON file.
+        single_category: Assessment category.
+        single_property: Assessment criteria.
+        single_value: Assessment value rule.
+    """
+    report_data = {
+        "Assessment Criteria": {
+            "Category": single_category,
+            "Property": single_property,
+            "Value": single_value,
+        },
+        "Results": data,
+    }
+    with open(filename, "w") as file:
+        json.dump(report_data, file, indent=4)
+
+
+def generate_pdf_report(
+    data: Dict[str, List[Dict[str, str]]],
+    filename: str,
+    single_category: str,
+    single_property: str,
+    single_value: str,
+) -> None:
+    """
+    Generates a PDF report. Suitable for official documentation.
+
+    Args:
+        data: Data to be included in the report.
+        filename: PDF file to save.
+        single_category: Assessment category.
+        single_property: Assessment criteria.
+        single_value: Assessment value rule.
+    """
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+    pdf.cell(200, 10, txt="Report", ln=True, align="C")
+    criteria_info = f"Criteria: {single_category} - {single_property} - {single_value}"
+    pdf.cell(200, 10, txt=criteria_info, ln=True)
+    pdf.cell(200, 10, txt="Name | Type | Family | ID | Status", ln=True)
+
+    for status, objects in data.items():
+        for obj in objects:
+            obj_info = f"{obj['name']} | {obj['type']} | {obj['family']} | {obj['id']} | {status}"
+            pdf.cell(200, 10, txt=obj_info, ln=True)
+
+    pdf.output(filename)
+
+
+def generate_html_report(
+    data: Dict[str, List[Dict[str, str]]],
+    single_category: str,
+    single_property: str,
+    single_value: str,
+) -> str:
+    """
+    Generates HTML content for the report. Easily styled and readable.
+
+    Args:
+        data: The data to display.
+        single_category: Assessment category.
+        single_property: Assessment criteria.
+        single_value: Assessment value rule.
+    """
+    html_content = "<html><head><title>Report</title></head><body>"
+    criteria_header = (
+        f"<h1>Report: {single_category} - {single_property} - {single_value}</h1>"
+    )
+    html_content += criteria_header
+    html_content += "<table border='1'>"
+    html_content += (
+        "<tr><th>Name</th><th>Type</th><th>Family</th><th>ID</th><th>Status</th></tr>"
+    )
+
+    for status, objects in data.items():
+        for obj in objects:
+            row = (
+                f"<tr><td>{obj['name']}</td><td>{obj['type']}</td>"
+                f"<td>{obj['family']}</td><td>{obj['id']}</td><td>{status}</td></tr>"
+            )
+            html_content += row
+
+    html_content += "</table></body></html>"
+    return html_content
+
+
+def generate_report(
+    assessed_objects: Dict[str, List[Dict[str, str]]],
+    report_format: str,
+    single_category: str,
+    single_property: str,
+    single_value: str,
+) -> str:
+    """
+    Main function to orchestrate report generation in various formats.
+
+    Args:
+        assessed_objects: Categorized assessment data.
+        report_format: The format to generate ('HTML', 'JSON', 'PDF').
+        single_category: Assessment category.
+        single_property: Assessment criteria.
+        single_value: Assessment value rule.
+    """
+    report_filename = f"report.{report_format.lower()}"
+
+    if report_format == "HTML":
+        html_report = generate_html_report(
+            assessed_objects, single_category, single_property, single_value
+        )
+        save_html_report(html_report, report_filename)
+    elif report_format == "JSON":
+        save_json_report(
+            assessed_objects,
+            report_filename,
+            single_category,
+            single_property,
+            single_value,
+        )
+    elif report_format == "PDF":
+        generate_pdf_report(
+            assessed_objects,
+            report_filename,
+            single_category,
+            single_property,
+            single_value,
+        )
+    else:
+        raise ValueError("Unsupported report format")
+
+    return report_filename

main.py

@@ -1,6 +1,7 @@
-"""This module contains the business logic of the function.
-
-Use the automation_context module to wrap your function in an Autamate context helper
+"""
+This module contains the business logic for a Speckle Automate function.
+It demonstrates how to define input models, traverse and process data,
+and generate reports based on user-specified criteria.
 """
 from enum import Enum
 
@@ -11,29 +12,44 @@ from speckle_automate import (
     execute_automate_function,
 )
 
-from flatten import flatten_base
+from Rules.checks import BaseObjectRules
+from Rules.traversal import get_data_traversal_rules
+from Utilities.report import generate_report
 
 
 class ThresholdMode(Enum):
+    """
+    ThresholdMode: Defines different modes for reporting thresholds.
+    """
+
     ERROR = "ERROR"
     WARN = "WARN"
     INFO = "INFO"
 
 
 class Format(Enum):
+    """
+    Format: Enum for defining report formats.
+    """
+
     PDF = "PDF"
     HTML = "HTML"
     JSON = "JSON"
 
 
 def create_one_of_enum(enum_cls):
+    """
+    Helper function to create a JSON schema from an Enum class.
+    This is used for generating user input forms in the UI.
+    """
     return [{"const": item.value, "title": item.name} for item in enum_cls]
 
 
 class FunctionInputs(AutomateBase):
-    """These are function author defined values.
+    """
+    FunctionInputs: Defines user inputs for the automation function.
+    The structure is based on Pydantic models for data validation.
 
-    Automate will make sure to supply them matching the types specified here.
     Please use the pydantic model schema to define your inputs:
     https://docs.pydantic.dev/latest/usage/models/
     """
@@ -95,50 +111,102 @@ def automate_function(
     automate_context: AutomationContext,
     function_inputs: FunctionInputs,
 ) -> None:
-    """This is an example Speckle Automate function.
+    """
+    The core logic of the Speckle Automate function.
+    Processes Speckle data and generates a report based on user inputs.
 
     Args:
-        automate_context: A context helper object, that carries relevant information
-            about the runtime context of this function.
-            It gives access to the Speckle project data, that triggered this run.
-            It also has convenience methods attach result data to the Speckle model.
-        function_inputs: An instance object matching the defined schema.
+        automate_context: Context object with data and methods for the run.
+        function_inputs: User-defined input values.
     """
     # the context provides a convenient way, to receive the triggering version
     version_root_object = automate_context.receive_version()
 
-    objects_with_forbidden_speckle_type = [
-        b
-        for b in flatten_base(version_root_object)
-        if b.speckle_type == function_inputs.forbidden_speckle_type
-    ]
-    count = len(objects_with_forbidden_speckle_type)
+    # Traverse the received Speckle data.
+    speckle_data = get_data_traversal_rules()
+    traversal_contexts_collection = speckle_data.traverse(version_root_object)
 
-    if count > 0:
-        # this is how a run is marked with a failure cause
-        automate_context.attach_error_to_objects(
-            category="Forbidden speckle_type"
-            " ({function_inputs.forbidden_speckle_type})",
-            object_ids=[o.id for o in objects_with_forbidden_speckle_type if o.id],
-            message="This project should not contain the type: "
-            f"{function_inputs.forbidden_speckle_type}",
-        )
-        automate_context.mark_run_failed(
-            "Automation failed: "
-            f"Found {count} object that have one of the forbidden speckle types: "
-            f"{function_inputs.forbidden_speckle_type}"
+    # Assuming each object has properties: name, type, and id
+    assessed_objects = {"missing": [], "invalid": [], "passing": []}
+
+    # Checking parameters
+    for context in traversal_contexts_collection:
+        current_object = getattr(context.current, "parameters", None)
+
+        if current_object:
+            for parameter_key in current_object.get_dynamic_member_names():
+                parameter = current_object[parameter_key]
+                assessment = BaseObjectRules.evaluate_parameter(
+                    parameter, function_inputs
+                )
+
+                if getattr(
+                    current_object, "speckle_type", None
+                ) == "Objects.Other.Revit.RevitInstance" and hasattr(
+                    current_object, "definition"
+                ):
+                    type_ = getattr(current_object.definition, "type", "Unknown")
+                    family = getattr(current_object.definition, "family", "Unknown")
+                else:
+                    type_ = getattr(current_object, "type", "Unknown")
+                    family = getattr(current_object, "family", "Unknown")
+
+                object_info = {
+                    "name": getattr(current_object, "name", "Unknown"),
+                    "type": type_,
+                    "family": family,
+                    "id": getattr(current_object, "id", "Unknown"),
+                }
+
+                if assessment:
+                    assessed_objects[assessment].append(object_info)
+
+    # Attach errors or info to objects based on their parameter evaluation state
+    for state, objects in assessed_objects.items():
+        ids = [obj["id"] for obj in objects if "id" in obj and obj["id"]]
+        if not ids:
+            continue
+
+        # Construct a detailed message for each object
+        detailed_messages = [
+            f"{obj['name']} (Type: {obj['type']}, ID: {obj['id']})"
+            for obj in objects
+            if "id" in obj and obj["id"]
+        ]
+
+        # Combine messages into a single string
+        combined_message = (
+            f"Found {len(objects)} objects with {state} parameters: "
+            + "; ".join(detailed_messages)
         )
 
-        # set the automation context view, to the original model / version view
-        # to show the offending objects
-        automate_context.set_context_view()
+        if state in ["missing", "invalid"]:
+            automate_context.attach_error_to_objects(
+                category=state.capitalize(), object_ids=ids, message=combined_message
+            )
+        else:  # 'valid'
+            automate_context.attach_info_to_objects(
+                category=state.capitalize(), object_ids=ids, message=combined_message
+            )
 
+    # Determine overall automation success or failure
+    if assessed_objects["missing"] or assessed_objects["invalid"]:
+        automate_context.mark_run_failed("Automation failed due to parameter issues.")
     else:
-        automate_context.mark_run_success("No forbidden types found.")
+        automate_context.mark_run_success("All parameters are valid.")
 
-    # if the function generates file results, this is how it can be
-    # attached to the Speckle project / model
-    # automate_context.store_file_result("./report.pdf")
+    # Generate and attach the report
+    report_format = (
+        function_inputs.report_format.value
+    )  # Accessing the value of the Enum
+    report_file = generate_report(
+        assessed_objects,
+        report_format,
+        function_inputs.single_category,
+        function_inputs.single_property,
+        function_inputs.single_rule,
+    )
+    automate_context.store_file_result(report_file)
 
 
 # make sure to call the function with the executor

poetry.lock

@@ -263,6 +263,16 @@ wrapt = ">=1.10,<2"
 [package.extras]
 dev = ["PyTest", "PyTest-Cov", "bump2version (<1)", "sphinx (<2)", "tox"]
 
+[[package]]
+name = "fpdf"
+version = "1.7.2"
+description = "Simple PDF generation for Python"
+optional = false
+python-versions = "*"
+files = [
+    {file = "fpdf-1.7.2.tar.gz", hash = "sha256:125840783289e7d12552b1e86ab692c37322e7a65b96a99e0ea86cca041b6779"},
+]
+
 [[package]]
 name = "gql"
 version = "3.4.1"
@@ -1178,4 +1188,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "c4fdb3557a519241ba9adbf905280be8542769843ba0178667fbbdcb6be57d73"
+content-hash = "62c33a6549afc140353a7eb6b0c7bc2b632d8feb1f4bb7d1fd616b3b030363d8"

pyproject.toml

@@ -8,6 +8,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.11"
 specklepy = "2.17.11"
+fpdf = "^1.7.2"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.3.0"

requirements.txt

@@ -110,6 +110,8 @@ charset-normalizer==3.3.2 ; python_version >= "3.11" and python_version < "4.0"
 deprecated==1.2.14 ; python_version >= "3.11" and python_version < "4.0" \
     --hash=sha256:6fac8b097794a90302bdbb17b9b815e732d3c4720583ff1b198499d78470466c \
     --hash=sha256:e5323eb936458dccc2582dc6f9c322c852a775a27065ff2b0c4970b9d53d01b3
+fpdf==1.7.2 ; python_version >= "3.11" and python_version < "4.0" \
+    --hash=sha256:125840783289e7d12552b1e86ab692c37322e7a65b96a99e0ea86cca041b6779
 gql[requests,websockets]==3.4.1 ; python_version >= "3.11" and python_version < "4.0" \
     --hash=sha256:11dc5d8715a827f2c2899593439a4f36449db4f0eafa5b1ea63948f8a2f8c545 \
     --hash=sha256:315624ca0f4d571ef149d455033ebd35e45c1a13f18a059596aeddcea99135cf
@@ -627,6 +629,7 @@ Pillow~=10.1.0
 docutils~=0.20.1
 sphinx~=7.0.1
 Jinja2~=3.1.2
+fpdf~=1.7.2
 mypy~=1.6.1
 filelock~=3.12.2
 idna~=3.4
@@ -652,3 +655,4 @@ stringcase~=1.2.0
 ujson~=5.8.0
 wrapt~=1.16.0
 httpx~=0.25.1
+python-dotenv~=1.0.0

tests/conftest.py

@@ -0,0 +1,24 @@
+import os
+
+from dotenv import load_dotenv
+
+
+def pytest_configure(config):
+    load_dotenv(dotenv_path=".env")
+
+    token_var = "SPECKLE_TOKEN"
+    server_var = "SPECKLE_SERVER_URL"
+    token = os.getenv(token_var)
+    server = os.getenv(server_var)
+
+    if not token:
+        raise ValueError(f"Cannot run tests without a {token_var} environment variable")
+
+    if not server:
+        raise ValueError(
+            f"Cannot run tests without a {server_var} environment variable"
+        )
+
+    # Set the token as an attribute on the config object
+    config.SPECKLE_TOKEN = token
+    config.SPECKLE_SERVER_URL = server

tests/test_function.py

@@ -1,19 +1,18 @@
 """Run integration tests with a speckle server."""
-import os
+
 import secrets
 import string
 
 import pytest
 from gql import gql
 from speckle_automate import (
+    AutomationContext,
     AutomationRunData,
     AutomationStatus,
     run_function,
 )
-from specklepy.api import operations
 from specklepy.api.client import SpeckleClient
 from specklepy.objects.base import Base
-from specklepy.transports.server import ServerTransport
 
 from main import FunctionInputs, automate_function
 
@@ -25,12 +24,12 @@ def crypto_random_string(length: int) -> str:
 
 
 def register_new_automation(
-    project_id: str,
-    model_id: str,
-    speckle_client: SpeckleClient,
-    automation_id: str,
-    automation_name: str,
-    automation_revision_id: str,
+        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(
@@ -67,19 +66,14 @@ def register_new_automation(
 
 
 @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
+def speckle_token(request) -> str:
+    return request.config.SPECKLE_TOKEN
 
 
 @pytest.fixture()
-def speckle_server_url() -> str:
+def speckle_server_url(request) -> 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")
+    return request.config.SPECKLE_SERVER_URL
 
 
 @pytest.fixture()
@@ -101,23 +95,18 @@ def test_object() -> Base:
 
 
 @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"
+# fixture to mock the AutomationRunData that would be generated by a full Automation Run
+def fake_automation_run_data(request, test_client: SpeckleClient) -> AutomationRunData:
+    SERVER_URL = request.config.SPECKLE_SERVER_URL
+    TOKEN = request.config.SPECKLE_TOKEN
 
-    model = test_client.branch.get(project_id, branch_name, commits_limit=1)
-    model_id: str = model.id
+    project_id = "9c6bfd2177"
+    model_id = "6193bdb540"
 
-    root_obj_id = operations.send(
-        test_object, [ServerTransport(project_id, test_client)]
-    )
-    version_id = test_client.commit.create(project_id, root_obj_id)
+    function_name = "Automate Density Check"
 
-    automation_name = crypto_random_string(10)
     automation_id = crypto_random_string(10)
+    automation_name = "Local Test Automation"
     automation_revision_id = crypto_random_string(10)
 
     register_new_automation(
@@ -129,30 +118,33 @@ def automation_run_data(
         automation_revision_id,
     )
 
-    automation_run_id = crypto_random_string(10)
-    function_id = crypto_random_string(10)
-    function_revision = crypto_random_string(10)
-    return AutomationRunData(
+    fake_run_data = AutomationRunData(
         project_id=project_id,
         model_id=model_id,
-        branch_name=branch_name,
-        version_id=version_id,
-        speckle_server_url=speckle_server_url,
+        branch_name="main",
+        version_id="107527ebd2",
+        speckle_server_url=SERVER_URL,
+        # These ids would be available with a valid registered Automation definition.
         automation_id=automation_id,
         automation_revision_id=automation_revision_id,
-        automation_run_id=automation_run_id,
-        function_id=function_id,
-        function_revision=function_revision,
+        automation_run_id=crypto_random_string(12),
+        # These ids would be available with a valid registered Function definition. Can also be faked.
+        function_id="12345",
+        function_name=function_name,
+        function_logo=None,
     )
 
+    return fake_run_data
 
-def test_function_run(automation_run_data: AutomationRunData, speckle_token: str):
+
+def test_function_run(fake_automation_run_data: AutomationRunData, speckle_token: str):
     """Run an integration test for the automate function."""
+    context = AutomationContext.initialize(fake_automation_run_data, speckle_token)
+
     automate_sdk = run_function(
+        context,
         automate_function,
-        automation_run_data,
-        speckle_token,
-        FunctionInputs(forbidden_speckle_type="Base"),
+        FunctionInputs(density_level=1000, max_percentage_high_density_objects=0.1),
     )
 
     assert automate_sdk.run_status == AutomationStatus.FAILED