diff --git a/tests/assets/plugins/blockifiers/blockifier_with_secrets.py b/tests/assets/plugins/blockifiers/blockifier_with_secrets.py new file mode 100644 index 000000000..2686f1238 --- /dev/null +++ b/tests/assets/plugins/blockifiers/blockifier_with_secrets.py @@ -0,0 +1,49 @@ +from typing import Type + +from pydantic import Field + +from steamship import Block, File +from steamship.invocable import Config, InvocableResponse, create_handler +from steamship.invocable.plugin_service import PluginRequest +from steamship.plugin.blockifier.blockifier import Blockifier +from steamship.plugin.inputs.raw_data_plugin_input import RawDataPluginInput +from steamship.plugin.outputs.block_and_tag_plugin_output import BlockAndTagPluginOutput + +# Note 1: this aligns with the same document in the internal Engine test. +# Note 2: This should be duplicated from the test_importer because of the way our test system will +# bundle this into an invocable deployment package (importing won't work!) +HANDLE = "test-blockifier-plugin-v1" +TEST_H1 = "A Poem" +TEST_S1 = "Roses are red." +TEST_S2 = "Violets are blue." +TEST_S3 = "Sugar is sweet, and I love you." +TEST_DOC = f"# {TEST_H1}\n\n{TEST_S1} {TEST_S2}\n\n{TEST_S3}\n" + + +class DummyBlockifierPlugin(Blockifier): + class DummyBlockifierConfig(Config): + secret: str = Field("") + + config: DummyBlockifierConfig + + @classmethod + def config_cls(cls) -> Type[Config]: + return cls.DummyBlockifierConfig + + def run( + self, request: PluginRequest[RawDataPluginInput] + ) -> InvocableResponse[BlockAndTagPluginOutput]: + return InvocableResponse( + data=BlockAndTagPluginOutput( + file=File( + blocks=[ + Block( + text=self.config.secret, + ), + ] + ) + ) + ) + + +handler = create_handler(DummyBlockifierPlugin) diff --git a/tests/steamship_tests/plugin/integration/test_e2e_blockifier_with_secrets.py b/tests/steamship_tests/plugin/integration/test_e2e_blockifier_with_secrets.py new file mode 100644 index 000000000..7e0bd3570 --- /dev/null +++ b/tests/steamship_tests/plugin/integration/test_e2e_blockifier_with_secrets.py @@ -0,0 +1,39 @@ +from steamship_tests import PLUGINS_PATH +from steamship_tests.utils.deployables import deploy_plugin +from steamship_tests.utils.fixtures import get_steamship_client + +from steamship import File + + +def test_e2e_blockifier_plugin(): + client = get_steamship_client() + blockifier_path = PLUGINS_PATH / "blockifiers" / "blockifier_with_secrets.py" + + # Send up the configTemplate that would be derived from this class: + # + # class DummyBlockifierConfig(Config): + # secret: str = Field("") + # + version_config_template = { + "secret": {"type": "string", "default": ""}, + } + + # Deploy with the secrets TOML {secret: "FOO"} + with deploy_plugin( + client, + blockifier_path, + "blockifier", + version_config_template=version_config_template, + secrets_toml='secret="FOO"', + ) as ( + plugin, + version, + instance, + ): + file = File.create(client=client, content="This is a test.") + assert len(file.refresh().blocks) == 0 + file.blockify(plugin_instance=instance.handle).wait() + file.refresh() + assert len(file.blocks) == 1 + assert file.blocks[0].text == "FOO" + file.delete() diff --git a/tests/steamship_tests/utils/deployables.py b/tests/steamship_tests/utils/deployables.py index 0e3b344be..154d31a7b 100644 --- a/tests/steamship_tests/utils/deployables.py +++ b/tests/steamship_tests/utils/deployables.py @@ -15,7 +15,11 @@ from steamship.data.user import User -def zip_deployable(file_path: Path, additional_requirements: Optional[List[str]] = None) -> bytes: +def zip_deployable( + file_path: Path, + additional_requirements: Optional[List[str]] = None, + secrets_toml: Optional[str] = None, +) -> bytes: """Prepare and zip a Steamship plugin.""" package_paths = [ @@ -42,6 +46,9 @@ def zip_deployable(file_path: Path, additional_requirements: Optional[List[str]] requirements_string += "\n" + "\n".join(additional_requirements) zip_file.writestr("requirements.txt", requirements_string) + if secrets_toml is not None: + zip_file.writestr(".steamship/secrets.toml", secrets_toml) + # Now we'll copy in the whole assets directory so that our test files can access things there.. for root, _, files in os.walk(TEST_ASSETS_PATH): for file in files: @@ -65,6 +72,7 @@ def deploy_plugin( safe_load_handler: bool = False, wait_for_init: bool = True, streaming: Optional[bool] = None, + secrets_toml: Optional[str] = None, ): plugin = Plugin.create( client, @@ -75,7 +83,7 @@ def deploy_plugin( is_public=False, ) - zip_bytes = zip_deployable(py_path) + zip_bytes = zip_deployable(py_path, secrets_toml=secrets_toml) hosting_handler = "steamship.invocable.entrypoint.safe_handler" if safe_load_handler else None plugin_version = PluginVersion.create( client,