fixup! fixup! fixup! fixup! fixup! fixup! fixup! Fix tests

snowflakedb · Jul 31, 2024 · 367a2f4 · 367a2f4
1 parent 70ae21b
commit 367a2f4
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/src/snowflake/cli/api/project/schemas/entities/streamlit_entity.py b/src/snowflake/cli/api/project/schemas/entities/streamlit_entity.py
@@ -19,10 +19,13 @@
 from pydantic import Field, model_validator
 from snowflake.cli.api.project.schemas.entities.common import EntityBase
 from snowflake.cli.api.project.schemas.identifier_model import ObjectIdentifierModel
+from snowflake.cli.api.project.schemas.updatable_model import (
+    DiscriminatorField,
+)
 
 
 class StreamlitEntity(EntityBase, ObjectIdentifierModel(object_name="Streamlit")):  # type: ignore
-    type: Literal["streamlit"]  # noqa: A003
+    type: Literal["streamlit"] = DiscriminatorField()  # noqa: A003
     title: Optional[str] = Field(
         title="Human-readable title for the Streamlit dashboard", default=None
     )

diff --git a/src/snowflake/cli/api/project/schemas/project_definition.py b/src/snowflake/cli/api/project/schemas/project_definition.py
@@ -172,12 +172,11 @@ def _validate_target_field(
         default=None,
     )
 
-
     def get_entities_by_type(self, entity_type: str):
         return {i: e for i, e in self.entities.items() if e.get_type() == entity_type}
 
 
-def build_project_definition(**data)-> ProjectDefinition:
+def build_project_definition(**data) -> ProjectDefinition:
     """
     Returns a ProjectDefinition instance with a version matching the provided definition_version value
     """

diff --git a/src/snowflake/cli/plugins/streamlit/commands.py b/src/snowflake/cli/plugins/streamlit/commands.py
@@ -134,8 +134,13 @@ def streamlit_deploy(
     stage is used. If the specified stage does not exist, the command creates it.
     """
 
-    pd = get_cli_context().project_definition
+    cli_context = get_cli_context()
+    pd = cli_context.project_definition
     if not pd.meets_version_requirement("2"):
+        if not pd.streamlit:
+            raise NoProjectDefinitionError(
+                project_type="streamlit", project_file=cli_context.project_root
+            )
         pd = _migrate_v1_streamlit_to_v2(pd)
 
     streamlits: Dict[str, StreamlitEntity] = pd.get_entities_by_type(
@@ -174,11 +179,6 @@ def streamlit_deploy(
 
 
 def _migrate_v1_streamlit_to_v2(pd: ProjectDefinition):
-    if not pd.streamlit:
-        raise NoProjectDefinitionError(
-            project_type="streamlit", project_file=cli_context.project_root
-        )
-
     default_env_file = "environment.yml"
     default_pages_dir = "pages"