Repository: great-expectations/great_expectations
Pull request #6987 (2023-01-30 15:07:56)

[MAINTENANCE] Add `update_expectation_suite` and `add_or_update_expectation_suite` to `DataContext` CRUD
Changes proposed in this pull request:
- Remainder of CRUD methods for suites
- Misc cleanup

### Definition of Done
- [x] My code follows the Great Expectations [style guide](https://docs.greatexpectations.io/docs/contributing/style_guides/code_style)
- [x] I have performed a [self-review](https://docs.greatexpectations.io/docs/contributing/contributing_checklist) of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have added [unit tests](https://docs.greatexpectations.io/docs/contributing/contributing_test#writing-unit-and-integration-tests) where applicable and made sure that new and existing tests are passing.
- [x] I have run any local integration tests and made sure that nothing is broken.


codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=6987&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations">Click on the image</a> for a new way to code review</b>
</summary>

  - Make big changes easier — review code in small groups of related files
  - Know where to start — see the whole change at a glance
  - Take a code tour — explore the change with an interactive tour
  - Make comments and review — all fully sync’ed with github

    [Try it now!](https://app.codesee.io/r/reviews?pr=6987&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations)
</details>

[![Review these changes using an interactive CodeSee Map](https://s3.us-east-2.amazonaws.com/maps.codesee.io/images/github/great-expectations/great_expectations/6987/4110dde9/94ed983c724a7285d3d9b148e752fd6a91dde936.svg)](https://app.codesee.io/r/reviews?pr=6987&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations)

#### Legend
![CodeSee Map Legend](https://app.codesee.io/Legend_v2.png)

Patch:
diff --git a/great_expectations/core/expectation_suite.py b/great_expectations/core/expectation_suite.py
index 5f4ab9795..d327db402 100644
--- a/great_expectations/core/expectation_suite.py
+++ b/great_expectations/core/expectation_suite.py
@@ -13,6 +13,7 @@ from typing import (
     Dict,
     List,
     Optional,
+    Sequence,
     Tuple,
     Type,
     Union,
@@ -94,7 +95,7 @@ class ExpectationSuite(SerializableDictDot):
         self,
         expectation_suite_name: str,
         data_context: Optional[AbstractDataContext] = None,
-        expectations: Optional[List[Union[dict, ExpectationConfiguration]]] = None,
+        expectations: Optional[Sequence[Union[dict, ExpectationConfiguration]]] = None,
         evaluation_parameters: Optional[dict] = None,
         data_asset_type: Optional[str] = None,
         execution_engine_type: Optional[Type[ExecutionEngine]] = None,
diff --git a/great_expectations/data_context/data_context/abstract_data_context.py b/great_expectations/data_context/data_context/abstract_data_context.py
index 31018c858..9b0fc59f3 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -766,8 +766,7 @@ class AbstractDataContext(ConfigPeer, ABC):
         datasource: Union[LegacyDatasource, BaseDatasource],
         save_changes: bool | None = None,
     ) -> None:
-        """
-        Updates a Datasource that already exists in the store.
+        """Updates a Datasource that already exists in the store.
 
         Args:
             datasource: The Datasource object to update.
@@ -812,8 +811,7 @@ class AbstractDataContext(ConfigPeer, ABC):
         name: str,
         **kwargs,
     ) -> LegacyDatasource | BaseDatasource:
-        """
-        Add a new Datasource or update an existing one on the context depending on whether
+        """Add a new Datasource or update an existing one on the context depending on whether
         it already exists or not. The configuration is provided as kwargs.
 
         Args:
@@ -2211,14 +2209,14 @@ class AbstractDataContext(ConfigPeer, ABC):
         self,
         expectation_suite_name: str,
         id: str | None = None,
-        expectations: list[dict | ExpectationConfiguration] | None = None,
+        expectations: Sequence[dict | ExpectationConfiguration] | None = None,
         evaluation_parameters: dict | None = None,
         data_asset_type: str | None = None,
         execution_engine_type: Type[ExecutionEngine] | None = None,
         meta: dict | None = None,
         overwrite_existing: bool = False,
         **kwargs,
-    ):
+    ) -> ExpectationSuite:
         if not isinstance(overwrite_existing, bool):
             raise ValueError("Parameter overwrite_existing must be of type BOOL")
 
@@ -2245,6 +2243,86 @@ class AbstractDataContext(ConfigPeer, ABC):
         self.expectations_store.set(key, expectation_suite, **kwargs)
         return expectation_suite
 
+    def update_expectation_suite(
+        self,
+        expectation_suite: ExpectationSuite,
+    ) -> None:
+        """Update an ExpectationSuite that already exists.
+
+        Args:
+            expectation_suite: The suite to use to update.
+
+        Raises:
+            DataContextError: A suite with the given name does not already exist.
+        """
+        expectation_suite_name: str = expectation_suite.expectation_suite_name
+        key = ExpectationSuiteIdentifier(expectation_suite_name=expectation_suite_name)
+        if not self.expectations_store.has_key(key):  # noqa: W601
+            raise gx_exceptions.DataContextError(
+                f"expectation_suite with name {expectation_suite_name} does not exist."
+            )
+
+        self._add_expectation_suite(
+            expectation_suite_name=expectation_suite_name,
+            id=expectation_suite.ge_cloud_id,
+            expectations=expectation_suite.expectations,
+            evaluation_parameters=expectation_suite.evaluation_parameters,
+            data_asset_type=expectation_suite.data_asset_type,
+            execution_engine_type=expectation_suite.execution_engine_type,
+            meta=expectation_suite.meta,  # type: ignore[has-type] # Should just be a dict but mypy has trouble inferring
+            overwrite_existing=True,
+        )
+
+    def add_or_update_expectation_suite(
+        self,
+        expectation_suite_name: str,
+        id: str | None = None,
+        expectations: list[dict | ExpectationConfiguration] | None = None,
+        evaluation_parameters: dict | None = None,
+        data_asset_type: str | None = None,
+        execution_engine_type: Type[ExecutionEngine] | None = None,
+        meta: dict | None = None,
+    ) -> ExpectationSuite:
+        """Add a new ExpectationSuite or update an existing one on the context depending on whether it already exists or not.
+
+        Args:
+            expectation_suite_name: The name of the suite to create.
+            id: Identifier to associate with this suite.
+            expectations: Expectation Configurations to associate with this suite.
+            evaluation_parameters: Evaluation parameters to be substituted when evaluating Expectations.
+            data_asset_type: Type of data asset to associate with this suite.
+            execution_engine_type: Name of the execution engine type.
+            meta: Metadata related to the suite.
+
+        Returns:
+            A new ExpectationSuite or an updated once (depending on whether or not it existed before this method call).
+        """
+        existing_suite = None
+        try:
+            existing_suite = self.get_expectation_suite(
+                expectation_suite_name=expectation_suite_name, ge_cloud_id=id
+            )
+        except gx_exceptions.DataContextError:
+            logger.info(
+                f"Could not find an existing suite named '{expectation_suite_name}'; creating a new one."
+            )
+
+        if existing_suite:
+            self.update_expectation_suite(existing_suite)
+            return self.get_expectation_suite(
+                expectation_suite_name=expectation_suite_name, ge_cloud_id=id
+            )
+
+        return self.add_expectation_suite(
+            expectation_suite_name=expectation_suite_name,
+            id=id,
+            expectations=expectations,
+            evaluation_parameters=evaluation_parameters,
+            data_asset_type=data_asset_type,
+            execution_engine_type=execution_engine_type,
+            meta=meta,
+        )
+
     def delete_expectation_suite(
         self,
         expectation_suite_name: Optional[str] = None,
@@ -2261,11 +2339,10 @@ class AbstractDataContext(ConfigPeer, ABC):
         key = ExpectationSuiteIdentifier(expectation_suite_name)  # type: ignore[arg-type]
         if not self.expectations_store.has_key(key):  # noqa: W601
             raise gx_exceptions.DataContextError(
-                "expectation_suite with name {} does not exist."
+                f"expectation_suite with name {expectation_suite_name} does not exist."
             )
-        else:
-            self.expectations_store.remove_key(key)
-            return True
+        self.expectations_store.remove_key(key)
+        return True
 
     @public_api
     @deprecated_argument(argument_name="ge_cloud_id", version="0.15.45")
diff --git a/great_expectations/profile/json_schema_profiler.py b/great_expectations/profile/json_schema_profiler.py
index 590302917..bb036ec8f 100644
--- a/great_expectations/profile/json_schema_profiler.py
+++ b/great_expectations/profile/json_schema_profiler.py
@@ -133,7 +133,7 @@ class JsonSchemaProfiler(Profiler):
                 }
             }
         suite = ExpectationSuite(
-            suite_name, expectations=expectations, meta=meta, data_context=None  # type: ignore[arg-type]
+            suite_name, expectations=expectations, meta=meta, data_context=None
         )
         suite.add_citation(
             comment=f"This suite was built by the {self.__class__.__name__}",


Test patch:
diff --git a/tests/data_context/test_data_context_state_management.py b/tests/data_context/test_data_context_state_management.py
index 1bc0201ba..3cd6785d6 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -10,7 +10,7 @@ from great_expectations.core.expectation_suite import ExpectationSuite
 from great_expectations.data_context.data_context.ephemeral_data_context import (
     EphemeralDataContext,
 )
-from great_expectations.data_context.store.store import Store
+from great_expectations.data_context.store.expectations_store import ExpectationsStore
 from great_expectations.data_context.types.base import (
     DataContextConfig,
     DatasourceConfig,
@@ -20,7 +20,7 @@ from great_expectations.data_context.types.base import (
 from great_expectations.exceptions.exceptions import StoreConfigurationError
 
 
-class StoreSpy(Store):
+class ExpectationsStoreSpy(ExpectationsStore):
     def __init__(self) -> None:
         self.save_count = 0
         super().__init__()
@@ -41,7 +41,7 @@ class EphemeralDataContextSpy(EphemeralDataContext):
     ) -> None:
         super().__init__(project_config)
         self.save_count = 0
-        self._expectations_store = StoreSpy()
+        self._expectations_store = ExpectationsStoreSpy()
 
     @property
     def expectations_store(self):
@@ -295,3 +295,89 @@ def test_add_expectation_suite_failure(
         in str(e.value)
     )
     assert context.expectations_store.save_count == 1
+
+
+@pytest.mark.unit
+def test_update_expectation_suite_success(
+    in_memory_data_context: EphemeralDataContextSpy,
+):
+    context = in_memory_data_context
+
+    suite_name = "default"
+    suite = context.add_expectation_suite(suite_name)
+
+    assert context.expectations_store.save_count == 1
+
+    suite.expectations = [
+        ExpectationConfiguration(
+            expectation_type="expect_column_values_to_be_in_set",
+            kwargs={"column": "x", "value_set": [1, 2, 4]},
+        ),
+    ]
+    context.update_expectation_suite(suite)
+
+    assert context.expectations_store.save_count == 2
+
+
+@pytest.mark.unit
+def test_update_expectation_suite_failure(
+    in_memory_data_context: EphemeralDataContextSpy,
+):
+    context = in_memory_data_context
+
+    suite_name = "my_brand_new_suite"
+    suite = ExpectationSuite(expectation_suite_name=suite_name)
+
+    with pytest.raises(gx_exceptions.DataContextError) as e:
+        context.update_expectation_suite(suite)
+
+    assert f"expectation_suite with name {suite_name} does not exist." in str(e.value)
+
+
+@pytest.mark.unit
+def test_add_or_update_expectation_suite_adds_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+):
+    context = in_memory_data_context
+
+    expectation_suite_name = "default"
+    expectations = [
+        ExpectationConfiguration(
+            expectation_type="expect_column_values_to_be_in_set",
+            kwargs={"column": "x", "value_set": [1, 2, 4]},
+        ),
+    ]
+    meta = {"great_expectations_version": "0.15.44"}
+
+    suite = context.add_or_update_expectation_suite(
+        expectation_suite_name=expectation_suite_name,
+        expectations=expectations,
+        meta=meta,
+    )
+
+    assert suite.expectation_suite_name == expectation_suite_name
+    assert suite.expectations == expectations
+    assert suite.meta == meta
+    assert context.expectations_store.save_count == 1
+
+
+@pytest.mark.unit
+def test_add_or_update_expectation_suite_updates_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+):
+    context = in_memory_data_context
+
+    suite_name = "default"
+    suite = context.add_expectation_suite(suite_name)
+
+    assert context.expectations_store.save_count == 1
+
+    suite.expectations = [
+        ExpectationConfiguration(
+            expectation_type="expect_column_values_to_be_in_set",
+            kwargs={"column": "x", "value_set": [1, 2, 4]},
+        ),
+    ]
+    context.add_or_update_expectation_suite(expectation_suite_name=suite_name)
+
+    assert context.expectations_store.save_count == 2