Floe Unit Tests

It is good practice to write unit tests for all floes. These tests can often quickly find trivial errors and improve development speed. The follow is an example of using the Artemis package to test floes.

Reader Writer Floe Test

from floe.api import WorkFloe
from orionplatform.cubes.datasets import DatasetReaderCube, DatasetWriterCube

floe = WorkFloe("Read Write Floe")

# Declare Cubes
ifs = DatasetReaderCube("ifs")
ofs = DatasetWriterCube("ofs")

# Add cubes to floe
floe.add_cubes(ifs, ofs)

# Promote parameters
ifs.promote_parameter(
    "data_in",
    promoted_name="in",
    title="Input Dataset of molecules",
)
ofs.promote_parameter(
    "data_out",
    promoted_name="out",
    title="Output File of Molecules",
)

# Connect Ports
ifs.success.connect(ofs.intake)
if __name__ == "__main__":
    floe.run()

# Testing previously defined floe
import os
from openeye import oechem
from artemis.test import FloeTestCase
from artemis.decorators import package
from artemis.wrappers import DatasetWrapper, WorkFloeWrapper, OutputDatasetWrapper

# Modify accordingly to get a relative path to the parent directory
PACKAGE_DIR = os.path.dirname(os.path.dirname(__file__))

FLOES_DIR = os.path.join(PACKAGE_DIR, "floes")

FILE_DIR = os.path.join(os.path.dirname(__file__), "test_data")

@package(PACKAGE_DIR)
class TestDatasetETL(FloeTestCase):
    mol_file = os.path.join(FILE_DIR, "molecules.smi")

    def test_reader_writer_floe(self):
        workfloe = WorkFloeWrapper.get_workfloe(
            os.path.join(FLOES_DIR, "readwrite.py"),
            run_timeout=600,
        )
        input_dataset = DatasetWrapper.from_file(self.mol_file)
        output_dataset = OutputDatasetWrapper(extension=".oedb")
        workfloe.start(
            {
                "promoted": {
                    "in": input_dataset.identifier,
                    "out": output_dataset.identifier,
                },
            }
        )
        self.assertWorkFloeComplete(workfloe)
        with open(output_dataset.path, "rb") as ifs:
            self.assertEqual(len(list(oechem.OEReadRecords(ifs))),100)