rename file so it works with ERs

client/verta/tests/deployable_entity/test_artifacts.py

@@ -5,6 +5,7 @@
 import hashlib
 import os
 import pickle
+import re
 import shutil
 import tempfile
 import zipfile
@@ -20,6 +21,7 @@
 )
 
 from .. import utils
+from ..registry.pydantic_models import InputClass, OutputClass
 
 
 def assert_model_packaging(deployable_entity, serialization, framework):
@@ -679,3 +681,33 @@ def test_setup_script(self, deployable_entity):
         assert deployable_entity.get_artifact(
             "setup_script"
         ).read() == six.ensure_binary(new_setup_script)
+
+
+class TestSetSchema:
+    def test_set_schema(self, deployable_entity):
+        input_schema = InputClass.schema()
+        output_schema = OutputClass.schema()
+        deployable_entity.set_schema(input=input_schema, output=output_schema)
+
+        assert deployable_entity.get_schema() == {
+            "input": input_schema,
+            "output": output_schema,
+        }
+
+    def test_set_schema_no_output(self, deployable_entity):
+        input_schema = InputClass.schema()
+        deployable_entity.set_schema(input=input_schema)
+
+        assert deployable_entity.get_schema() == {
+            "input": input_schema,
+        }
+
+    def test_set_schema_no_input_which_is_bad(self, deployable_entity):
+        output_schema = OutputClass.schema()
+        with pytest.raises(
+            TypeError,
+            match=re.escape(
+                "_DeployableEntity.set_schema() missing 1 required positional argument: 'input'"
+            ),
+        ):
+            deployable_entity.set_schema(output=output_schema)

client/verta/tests/registry/conftest.py

@@ -22,7 +22,7 @@ def docker_image():
 
 @pytest.fixture
 def make_model_schema_file(tmp_path, monkeypatch):
-    path = tmp_path / "model_schema.json"
+    path = tmp_path / "model_schema_json"
     schema = {"input": InputClass.schema(), "output": OutputClass.schema()}
     path.write_text(json.dumps(schema))
     monkeypatch.setenv(_MODEL_SCHEMA_PATH_ENV_VAR, str(path))

client/verta/verta/registry/_validate_schema.py

@@ -80,7 +80,7 @@ def main():
     def wrapper(self, input: Dict):
         # fetch schema
         model_schema_path = os.environ.get(
-            _MODEL_SCHEMA_PATH_ENV_VAR, "/app/model_schema.json"
+            _MODEL_SCHEMA_PATH_ENV_VAR, "/app/model_schema_json"
         )
         try:
             with open(model_schema_path, "r") as file:

client/verta/verta/tracking/entities/_deployable_entity.py

@@ -129,7 +129,7 @@ def set_schema(self, input: dict, output: dict = None) -> None:
             temp_filename = temp_file.name
             with open(temp_filename, 'w') as file:
                 json.dump(schema, file)
-            self.log_artifact(key="model_schema.json", artifact=temp_filename, overwrite=True)
+            self.log_artifact(key="model_schema_json", artifact=temp_filename, overwrite=True)
 
     def get_schema(self) -> Dict[str, dict]:
         """
@@ -149,7 +149,7 @@ def get_schema(self) -> Dict[str, dict]:
             Input and output json schemas.
 
         """
-        schema = self.get_artifact("model_schema.json")
+        schema = self.get_artifact("model_schema_json")
         return json.load(schema)
 
     @abc.abstractmethod