fix unit tests and formatting

Nike-Inc · Nov 7, 2023 · 44e4e9e · 44e4e9e
1 parent 131625e
commit 44e4e9e
Show file tree

Hide file tree

Showing 9 changed files with 19 additions and 16 deletions.
diff --git a/brickflow/engine/task.py b/brickflow/engine/task.py
@@ -32,7 +32,11 @@
  BrickflowProjectDeploymentSettings,
  get_brickflow_version,
 )
-from brickflow.bundles.model import JobsTasksNotebookTask, JobsTasksNotificationSettings, JobsTasksHealthRules
+from brickflow.bundles.model import (
+ JobsTasksNotebookTask,
+ JobsTasksNotificationSettings,
+ JobsTasksHealthRules,
+)
 from brickflow.cli.projects import DEFAULT_BRICKFLOW_VERSION_MODE
 from brickflow.context import (
  BrickflowBuiltInTaskVariables,
@@ -480,7 +484,6 @@ class Task:
  custom_execute_callback: Optional[Callable] = None
  ensure_brickflow_plugins: bool = False
  health: Optional[List[JobsTasksHealthRules]] = None
-
 
  def __post_init__(self) -> None:
  self.is_valid_task_signature()

diff --git a/brickflow/engine/workflow.py b/brickflow/engine/workflow.py
@@ -11,7 +11,7 @@
  JobsWebhookNotifications,
  JobsNotificationSettings,
  JobsTrigger,
- JobsHealthRules
+ JobsHealthRules,
 )
 from brickflow.context import BrickflowInternalVariables
 from brickflow.engine import ROOT_NODE
@@ -117,6 +117,7 @@ class Workflow:
  schedule_pause_status: str = "UNPAUSED"
  default_cluster: Optional[Cluster] = None
  clusters: List[Cluster] = field(default_factory=lambda: [])
+
  health: Optional[List[JobsHealthRules]] = None
  default_task_settings: TaskSettings = TaskSettings()
  email_notifications: Optional[WorkflowEmailNotifications] = None

diff --git a/tests/codegen/expected_bundles/dev_bundle_monorepo.yml b/tests/codegen/expected_bundles/dev_bundle_monorepo.yml
@@ -8,6 +8,7 @@ environments:
  email_notifications: null
  notification_settings: null
  webhook_notifications: null
+ health: null
  git_source:
  git_commit: a
  git_provider: github

diff --git a/tests/codegen/expected_bundles/dev_bundle_polyrepo.yml b/tests/codegen/expected_bundles/dev_bundle_polyrepo.yml
@@ -8,6 +8,7 @@ environments:
  email_notifications: null
  notification_settings: null
  webhook_notifications: null
+ health: null
  git_source:
  git_commit: a
  git_provider: github

diff --git a/tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml b/tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml
@@ -6,6 +6,7 @@ environments:
  jobs:
  some_wf:
  email_notifications: null
+ health: null
  git_source:
  git_commit: a
  git_provider: github
@@ -64,6 +65,7 @@ environments:
  email_notifications: null
  notification_settings: null
  webhook_notifications: null
+ health: null
  git_source:
  git_commit: a
  git_provider: github

diff --git a/tests/codegen/expected_bundles/local_bundle.yml b/tests/codegen/expected_bundles/local_bundle.yml
@@ -8,6 +8,7 @@ environments:
  email_notifications: null
  notification_settings: null
  webhook_notifications: null
+ health: null
  git_source: null
  job_clusters: []
  max_concurrent_runs: 1.0

diff --git a/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml b/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml
@@ -8,6 +8,7 @@ environments:
  email_notifications: null
  notification_settings: null
  webhook_notifications: null
+ health: null
  git_source: null
  job_clusters: []
  max_concurrent_runs: 1.0

diff --git a/tests/engine/sample_workflow.py b/tests/engine/sample_workflow.py
@@ -19,15 +19,11 @@
  ),
  tags={"test": "test2"},
  common_task_parameters={"all_tasks1": "test", "all_tasks3": "123"}, # type: ignore
- health = {
+ health={
  "rules": [
- {
- "metric": "RUN_DURATION_SECONDS",
- "op": "GREATER_THAN",
- "value": 7200
- } 
+ {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200}
  ]
- }
+ },
 )
 
 

diff --git a/tests/engine/test_workflow.py b/tests/engine/test_workflow.py
@@ -182,14 +182,11 @@ def test_tags(self):
  def test_default_task_settings(self):
  assert wf.default_task_settings is not None
 
- def test_health_settings(self): 
+ def test_health_settings(self):
  assert wf.health == {
  "rules": [
- {
- "metric": "RUN_DURATION_SECONDS",
- "op": "GREATER_THAN",
- "value": 7200
- } ]
+ {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200}
+ ]
  }
 
  def test_user(self):