Commit 4d8e42cb authored by ozzeh's avatar ozzeh
Browse files

Blacken

parent aae2a364
from .client import SchemaRegistry
from .errors import SchemaRegistryError, ModelNotRegisteredError
\ No newline at end of file
from .errors import SchemaRegistryError, ModelNotRegisteredError
......@@ -49,7 +49,9 @@ class Schema:
version.schema_version
] = self._get_schema_version_content(version)
def _get_schema_version_content(self, schema_version: _SchemaVersionModel) -> _SchemaContentModel:
def _get_schema_version_content(
self, schema_version: _SchemaVersionModel
) -> _SchemaContentModel:
describe_opts = dict(
RegistryName=self.registry_name,
SchemaName=self.schema_name,
......@@ -127,9 +129,11 @@ class SchemaRegistry:
schema_info = _SchemaCreateUpdateModel.parse_obj(response)
self._model_schemas[model] = schema_info
return schema_info
def schema_for_model(self, model: Type[BaseModel]) -> dict:
if model not in self._model_schemas:
raise ModelNotRegisteredError(model)
return self._model_schemas[model].dict(include=set(["schema_arn", "schema_name", "schema_version"]))
return self._model_schemas[model].dict(
include=set(["schema_arn", "schema_name", "schema_version"])
)
class SchemaRegistryError(Exception):
pass
class ModelNotRegisteredError(SchemaRegistryError):
def __init__(self, model):
self.model = model
\ No newline at end of file
self.model = model
......@@ -20,6 +20,7 @@ def named_registry():
_registry = SchemaRegistry(registry_name="TAPI-TEST")
yield _registry
@pytest.fixture(scope="module")
def dynamic_registry():
_registry = SchemaRegistry()
......@@ -43,6 +44,7 @@ def complex_model():
class ComplexModel(BaseModel):
"""Hi mom"""
name: str
description: Optional[str]
groups: List[Group]
......@@ -52,15 +54,15 @@ def complex_model():
def test_load_named_schemas(test_model, named_registry):
pass
def test_schema_registration(test_model, named_registry):
schema_info = named_registry.register_model("com.pleaseignore.tvm.test", test_model)
def test_registered_model(test_model, named_registry):
model_info = named_registry.schema_for_model(test_model)
def test_complex_mode(complex_model, named_registry):
named_registry.register_model("com.pleaseignore.tvm.test", complex_model)
\ No newline at end of file
named_registry.register_model("com.pleaseignore.tvm.test", complex_model)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment