Repository: great-expectations/great_expectations
Pull request #6926 (2023-01-25 19:07:47)

[MAINTENANCE] Add `add_expectation_suite` to `DataContext` CRUD
Changes proposed in this pull request:
- Add `add_expectation_suite` and related tests
- 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.


netlify[bot]: ### <span aria-hidden="true">👷</span> Deploy Preview for *niobium-lead-7998* processing.


|  Name | Link |
|---------------------------------|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | 7a5960b7c2dcc86218a9397963b468276a7c5495 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/niobium-lead-7998/deploys/63d17e05347e430008e1ffea |
codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=6926&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=6926&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/6926/4c905d6d/7a5960b7c2dcc86218a9397963b468276a7c5495.svg)](https://app.codesee.io/r/reviews?pr=6926&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/data_context/data_context/abstract_data_context.py b/great_expectations/data_context/data_context/abstract_data_context.py
index 4862b5dca..31018c858 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -22,6 +22,7 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -52,6 +53,7 @@ from great_expectations.core.config_provider import (
     _EnvironmentConfigurationProvider,
     _RuntimeEnvironmentConfigurationProvider,
 )
+from great_expectations.core.expectation_configuration import ExpectationConfiguration
 from great_expectations.core.expectation_validation_result import get_metric_kwargs_id
 from great_expectations.core.id_dict import BatchKwargs
 from great_expectations.core.run_identifier import RunIdentifier
@@ -2140,11 +2142,95 @@ class AbstractDataContext(ConfigPeer, ABC):
             ValueError: The input `overwrite_existing` is of the wrong type.
             DataContextError: A suite with the same name already exists (and `overwrite_existing` is not enabled).
         """
+        return self._add_expectation_suite(
+            expectation_suite_name=expectation_suite_name,
+            overwrite_existing=overwrite_existing,
+        )
+
+    def add_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:
+        """Build a new ExpectationSuite and save it utilizing the context's underlying ExpectationsStore.
+
+        Note that this method can be called by itself or run within the get_validator workflow.
+
+        When run with create_expectation_suite()::
+
+            expectation_suite_name = "genres_movies.fkey"
+            context.create_expectation_suite(expectation_suite_name, overwrite_existing=True)
+            batch = context.get_batch(
+                expectation_suite_name=expectation_suite_name
+            )
+
+
+        When run as part of get_validator()::
+
+            validator = context.get_validator(
+                datasource_name="my_datasource",
+                data_connector_name="whole_table",
+                data_asset_name="my_table",
+                create_expectation_suite_with_name="my_expectation_suite",
+            )
+            validator.expect_column_values_to_be_in_set("c1", [4,5,6])
+
+
+        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 built with provided input args.
+
+        Raises:
+            DataContextError: A suite with the same name already exists (and `overwrite_existing` is not enabled).
+        """
+        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,
+            overwrite_existing=False,  # `add` does not resolve collisions
+        )
+
+    def _add_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,
+        overwrite_existing: bool = False,
+        **kwargs,
+    ):
         if not isinstance(overwrite_existing, bool):
             raise ValueError("Parameter overwrite_existing must be of type BOOL")
 
         expectation_suite = ExpectationSuite(
-            expectation_suite_name=expectation_suite_name, data_context=self
+            expectation_suite_name=expectation_suite_name,
+            data_context=self,
+            ge_cloud_id=id,
+            expectations=expectations,
+            evaluation_parameters=evaluation_parameters,
+            data_asset_type=data_asset_type,
+            execution_engine_type=execution_engine_type,
+            meta=meta,
         )
         key = ExpectationSuiteIdentifier(expectation_suite_name=expectation_suite_name)
         if (
@@ -2152,10 +2238,9 @@ class AbstractDataContext(ConfigPeer, ABC):
             and not overwrite_existing
         ):
             raise gx_exceptions.DataContextError(
-                "expectation_suite with name {} already exists. If you would like to overwrite this "
-                "expectation_suite, set overwrite_existing=True.".format(
-                    expectation_suite_name
-                )
+                f"expectation_suite with name {expectation_suite_name} already exists."
+                " If you would like to overwrite this expectation_suite, please delete or"
+                " update it using `delete_expectation_suite` or `update_expectation_suite`, respectively."
             )
         self.expectations_store.set(key, expectation_suite, **kwargs)
         return expectation_suite


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 3a6ce5ffb..1bc0201ba 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -4,9 +4,13 @@ from typing import Mapping
 
 import pytest
 
+import great_expectations.exceptions as gx_exceptions
+from great_expectations.core.expectation_configuration import ExpectationConfiguration
+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.types.base import (
     DataContextConfig,
     DatasourceConfig,
@@ -16,6 +20,16 @@ from great_expectations.data_context.types.base import (
 from great_expectations.exceptions.exceptions import StoreConfigurationError
 
 
+class StoreSpy(Store):
+    def __init__(self) -> None:
+        self.save_count = 0
+        super().__init__()
+
+    def set(self, key, value, **kwargs):
+        self.save_count += 1
+        return super().set(key=key, value=value, **kwargs)
+
+
 class EphemeralDataContextSpy(EphemeralDataContext):
     """
     Simply wraps around EphemeralDataContext but keeps tabs on specific method calls around state management.
@@ -27,6 +41,11 @@ class EphemeralDataContextSpy(EphemeralDataContext):
     ) -> None:
         super().__init__(project_config)
         self.save_count = 0
+        self._expectations_store = StoreSpy()
+
+    @property
+    def expectations_store(self):
+        return self._expectations_store
 
     def _save_project_config(self):
         """
@@ -209,3 +228,70 @@ def test_add_or_update_datasource_adds_successfully(
     assert num_datasource_after == num_datasource_before + 1
     assert num_datasource_configs_after == num_datasource_configs_before + 1
     assert context.save_count == 1
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "kwargs,expected_suite",
+    [
+        pytest.param(
+            {"expectation_suite_name": "my_new_suite"},
+            ExpectationSuite(expectation_suite_name="my_new_suite"),
+            id="only name",
+        ),
+        pytest.param(
+            {
+                "expectations": [
+                    ExpectationConfiguration(
+                        expectation_type="expect_column_values_to_be_in_set",
+                        kwargs={"column": "x", "value_set": [1, 2, 4]},
+                    ),
+                ],
+                "expectation_suite_name": "default",
+                "meta": {"great_expectations_version": "0.15.44"},
+            },
+            ExpectationSuite(
+                expectations=[
+                    ExpectationConfiguration(
+                        expectation_type="expect_column_values_to_be_in_set",
+                        kwargs={"column": "x", "value_set": [1, 2, 4]},
+                    ),
+                ],
+                expectation_suite_name="default",
+                meta={"great_expectations_version": "0.15.44"},
+            ),
+            id="misc args",
+        ),
+    ],
+)
+def test_add_expectation_suite_success(
+    in_memory_data_context: EphemeralDataContextSpy,
+    kwargs: dict,
+    expected_suite: ExpectationSuite,
+):
+    context = in_memory_data_context
+
+    suite = context.add_expectation_suite(**kwargs)
+
+    assert suite == expected_suite
+    assert context.expectations_store.save_count == 1
+
+
+@pytest.mark.unit
+def test_add_expectation_suite_failure(
+    in_memory_data_context: EphemeralDataContextSpy,
+):
+    context = in_memory_data_context
+
+    suite_name = "default"
+    context.add_expectation_suite(expectation_suite_name=suite_name)
+
+    with pytest.raises(gx_exceptions.DataContextError) as e:
+        context.add_expectation_suite(expectation_suite_name=suite_name)
+
+    assert f"expectation_suite with name {suite_name} already exists" in str(e.value)
+    assert (
+        "please delete or update it using `delete_expectation_suite` or `update_expectation_suite`"
+        in str(e.value)
+    )
+    assert context.expectations_store.save_count == 1