Artemis Examples

Handling Local and Orion Floe testing

Artemis can run tests either locally or upload packages to Orion and run the floes on Orion. Sometimes it is necessary to distinguish between running tests locally rather than on Orion. The following example shows how to use the using_orion() to do this.

Note

Differs from orionclient.session.in_orion() in that using_orion() indicates that the tests are running against Orion using local credentials, while orionclient.session.in_orion() indicates that the code is being run from within Orion.

import os
from time import sleep

# Python Floe Package pip installed with -e
import classic_floes

from artemis.test import FloeTestCase
from artemis.decorators import package
from artemis.constants import WORKFLOE_COMPLETE
from artemis.wrappers import (
    DatasetWrapper,
    WorkFloeWrapper,
    OutputDatasetWrapper,
    using_orion,
)

# Path to the top level directory
PACKAGE_DIR = os.path.dirname(os.path.dirname(classic_floes.__file__))

FILE_DIR = os.path.join(PACKAGE_DIR, "tests", "data")
FLOE_DIR = os.path.join(PACKAGE_DIR, "floes")


@package(PACKAGE_DIR)
class ExampleFloeTest(FloeTestCase):
    def test_example_floe(self):
        workfloe = WorkFloeWrapper.get_workfloe(
            os.path.join(FLOE_DIR, "classic_omega.py"),  # Path to the local workfloe
            run_timeout=1200,  # Max time to run
            queue_timeout=600,  # Max time to spend queued
        )
        input_dataset = DatasetWrapper.from_file(
            os.path.join(FILE_DIR, "eMol_ran50.ism")
        )
        output_dataset = OutputDatasetWrapper()
        failure_dataset = OutputDatasetWrapper()
        parameters = {
            "promoted": {
                "Input Molecules": input_dataset.identifier,
                "out": output_dataset.identifier,
                "failed": failure_dataset.identifier,
                "Maximum Conformers": 10,
            }
        }
        if using_orion():
            # Floe will be running in Orion, and can be run non-blocking
            workfloe.start(parameters, block=False)
            while workfloe.state != WORKFLOE_COMPLETE:
                print(f"Workfloe running, state: {workfloe.state}")
                sleep(5)
        else:
            workfloe.start(parameters)
        self.assertWorkFloeComplete(workfloe)

Testing with Collections

Artemis provides wrapper that can be used to construct and retrieve shard collections when interacting with Floes. The following code demonstrates using the CollectionWrapper and OutputCollectionWrapper.

import os

# Python Floe Package pip installed with -e
import my_package

from artemis.test import FloeTestCase
from artemis.decorators import package
from artemis.wrappers import WorkFloeWrapper, CollectionWrapper, OutputCollectionWrapper

# Path to the top level directory
PACKAGE_DIR = os.path.dirname(os.path.dirname(my_package.__file__))

# Should conform to my package
FILE_DIR = os.path.join(PACKAGE_DIR, "tests", "data")
FLOE_DIR = os.path.join(PACKAGE_DIR, "floes")


@package(PACKAGE_DIR)
class ExampleFloeTest(FloeTestCase):
    def test_collection_floe(self):
        workfloe = WorkFloeWrapper.get_workfloe(
            os.path.join(FLOE_DIR, "collection_floe.py"),  # Path to the local workfloe
            run_timeout=1200,  # Max time to run
            queue_timeout=600,  # Max time to spend queued
        )
        # Will split smiles into records and split into 50 shards
        input_collection = CollectionWrapper.from_file(
            os.path.join(FILE_DIR, "100.ism"), num_shards=50
        )
        output_collection = OutputCollectionWrapper()
        parameters = {
            "promoted": {
                "in": input_collection.identifier,
                "out": output_collection.identifier,
            }
        }
        workfloe.start(parameters)
        self.assertWorkFloeComplete(workfloe)
        result_collection = output_collection.get()
        for shard in result_collection.list_shards():
            print(f"Shard {shard.id}")

Dynamic Packaging

Dynamic packaging is useful when you want to write test floes, but don’t want to provide them to users along with your package. By configuring a package dynamically you can also modify the requirements without making permanent changes to the package.

from tempfile import NamedTemporaryFile

from artemis.test import FloeTestCase
from artemis.decorators import package
from artemis.packaging import OrionTestPackage

test_package = OrionTestPackage(manifest=dict(requirements="custom_requirements.txt"))
temp_req = NamedTemporaryFile(suffix=".txt")
with open(temp_req.name, "w") as ofs:
    # Create a file with custom requirements
    ofs.write("openeye-orionplatform==3.1.0\\n")

# Add the contents of the regular package
test_package.add_directory("path/to/package/directory")

# Remove unnecssary directories or files
test_package.remove_directory("package/package_data/")
test_package.remove_file("requirements.txt")
# Add in files and directories to the package
test_package.add_file(temp_req.name, dest="custom_requirements.txt")


# Package the package to the package decorator
@package(test_package)
class BasicDynamicPackageTestCase(FloeTestCase):
    def test_dynamic_package(self):
        # The package will have built once the test runs
        pass