content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  import os
  import random
  import string
  from contextlib import ExitStack as does_not_raise
  from contextlib import contextmanager
  from typing import Any, Dict, Generator, Optional
  from uuid import uuid4
  
  import pytest
  from pydantic import BaseModel
  from typing_extensions import Annotated
  
  from tests.integration.functional.conftest import (
      constant_int_output_test_step,
      int_plus_one_test_step,
  )
  from tests.integration.functional.utils import sample_name
  from zenml import (
      ExternalArtifact,
      get_pipeline_context,
      get_step_context,
      log_metadata,
      pipeline,
      save_artifact,
      step,
  )
  from zenml.client import Client
  from zenml.config.source import Source
  from zenml.constants import PAGE_SIZE_DEFAULT
  from zenml.enums import (
      ArtifactSaveType,
      MetadataResourceTypes,
      ModelStages,
      StackComponentType,
      TaggableResourceTypes,
  )
  from zenml.exceptions import (
      EntityExistsError,
      IllegalOperationError,
      InitializationException,
  )
  from zenml.io import fileio
  from zenml.model.model import Model
  from zenml.models import (
      ComponentResponse,
      ModelResponse,
      ModelVersionResponse,
      PipelineBuildRequest,
      PipelineRequest,
      PipelineSnapshotRequest,
      RunMetadataResource,
      StackResponse,
      TagResource,
  )
  from zenml.utils import io_utils
  from zenml.utils.string_utils import random_str
  
  
  def _create_local_orchestrator(
      client: Client,
      orchestrator_name: str = "OrchesTraitor",
  ) -> ComponentResponse:
      return client.create_stack_component(
          name=orchestrator_name,
          flavor="local",
          component_type=StackComponentType.ORCHESTRATOR,
          configuration={},
      )
  
  
  def _create_local_artifact_store(
      client: Client,
      artifact_store_name: str = "Art-E-Fact",
  ) -> ComponentResponse:
      return client.create_stack_component(
          name=artifact_store_name,
          flavor="local",
          component_type=StackComponentType.ARTIFACT_STORE,
          configuration={},
      )
  
  
  def _create_local_stack(
      client: Client,
      stack_name: str,
      orchestrator_name: Optional[str] = None,
      artifact_store_name: Optional[str] = None,
  ) -> StackResponse:
      """Creates a local stack with components with the given names. If the names are not given, a random string is used instead."""
  
      def _random_name():
          return "".join(random.choices(string.ascii_letters, k=10))
  
      orchestrator = _create_local_orchestrator(
          client=client, orchestrator_name=orchestrator_name or _random_name()
      )
  
      artifact_store = _create_local_artifact_store(
          client=client,
          artifact_store_name=artifact_store_name or _random_name(),
      )
  
      return client.create_stack(
          name=stack_name,
          components={
              StackComponentType.ORCHESTRATOR: str(orchestrator.id),
              StackComponentType.ARTIFACT_STORE: str(artifact_store.id),
          },
      )
  
  
  def test_repository_detection(tmp_path):
      """Tests detection of ZenML repositories in a directory."""
      assert Client.is_repository_directory(tmp_path) is False
      Client.initialize(tmp_path)
      assert Client.is_repository_directory(tmp_path) is True
  
  
  def test_initializing_repo_creates_directory_and_uses_default_stack(
      tmp_path, clean_client
  ):
      """Tests that repo initialization creates a .zen directory and uses the default local stack."""
      Client.initialize(tmp_path)
      assert fileio.exists(str(tmp_path / ".zen"))
  
      client = Client()
      # switch to the new repo root
      client.activate_root(tmp_path)
  
      stack = client.active_stack_model
      assert isinstance(
          stack.components[StackComponentType.ORCHESTRATOR][0],
          ComponentResponse,
      )
      assert isinstance(
          stack.components[StackComponentType.ARTIFACT_STORE][0],
          ComponentResponse,
      )
      with pytest.raises(KeyError):
          assert stack.components[StackComponentType.CONTAINER_REGISTRY]
  
  
  def test_initializing_repo_twice_fails(tmp_path):
      """Tests that initializing a repo in a directory where another repo already exists fails."""
      Client.initialize(tmp_path)
      with pytest.raises(InitializationException):
          Client.initialize(tmp_path)
  
  
  def test_freshly_initialized_repo_attributes(tmp_path):
      """Tests that the attributes of a new repository are set correctly."""
      Client.initialize(tmp_path)
      client = Client(tmp_path)
  
      assert client.root == tmp_path
  
  
  def test_finding_repository_directory_with_explicit_path(
      tmp_path, clean_client
  ):
      """Tests that a repository can be found using an explicit path, an environment variable and the current working directory."""
      subdirectory_path = tmp_path / "some_other_directory"
      io_utils.create_dir_recursive_if_not_exists(str(subdirectory_path))
      os.chdir(str(subdirectory_path))
  
      # no repo exists and explicit path passed
      assert Client.find_repository(tmp_path) is None
      assert Client(tmp_path).root is None
  
      # no repo exists and no path passed (=uses current working directory)
      assert Client.find_repository() is None
      Client._reset_instance()
      assert Client().root is None
  
      # no repo exists and explicit path set via environment variable
      os.environ["ZENML_REPOSITORY_PATH"] = str(tmp_path)
      assert Client.find_repository() is None
      Client._reset_instance()
      assert Client().root is None
  
      del os.environ["ZENML_REPOSITORY_PATH"]
  
      # initializing the repo
      Client.initialize(tmp_path)
  
      # repo exists and explicit path passed
      assert Client.find_repository(tmp_path) == tmp_path
      assert Client(tmp_path).root == tmp_path
  
      # repo exists and explicit path to subdirectory passed
      assert Client.find_repository(subdirectory_path) is None
      assert Client(subdirectory_path).root is None
  
      # repo exists and no path passed (=uses current working directory)
      assert Client.find_repository() == tmp_path
      Client._reset_instance()
      assert Client().root == tmp_path
  
      # repo exists and explicit path set via environment variable
      os.environ["ZENML_REPOSITORY_PATH"] = str(tmp_path)
      assert Client.find_repository() == tmp_path
      Client._reset_instance()
      assert Client().root == tmp_path
  
      # repo exists and explicit path to subdirectory set via environment variable
      os.environ["ZENML_REPOSITORY_PATH"] = str(subdirectory_path)
      assert Client.find_repository() is None
      Client._reset_instance()
      assert Client().root is None
  
      del os.environ["ZENML_REPOSITORY_PATH"]
  
  
  def test_activating_nonexisting_stack_fails(clean_client):
      """Tests that activating a stack name that isn't registered fails."""
      with pytest.raises(KeyError):
          clean_client.activate_stack(str(uuid4()))
  
  
  def test_activating_a_stack_updates_the_config_file(clean_client):
      """Tests that the newly active stack name gets persisted."""
      stack = _create_local_stack(client=clean_client, stack_name="new_stack")
      clean_client.activate_stack(stack.id)
  
      assert Client(clean_client.root).active_stack_model.name == stack.name
  
  
  def test_registering_a_stack(clean_client):
      """Tests that registering a stack works and the stack gets persisted."""
      orch = _create_local_orchestrator(
          client=clean_client,
      )
      art = _create_local_artifact_store(
          client=clean_client,
      )
      new_stack_name = "some_new_stack_name"
      new_stack = clean_client.create_stack(
          name=new_stack_name,
          components={
              StackComponentType.ORCHESTRATOR: str(orch.id),
              StackComponentType.ARTIFACT_STORE: str(art.id),
          },
      )
  
      Client(clean_client.root)
      with does_not_raise():
          clean_client.zen_store.get_stack(new_stack.id)
  
  
  def test_registering_a_stack_with_existing_name(clean_client):
      """Tests that registering a stack for an existing name fails."""
      _create_local_stack(
          client=clean_client,
          stack_name="axels_super_awesome_stack_of_fluffyness",
      )
      orchestrator = _create_local_orchestrator(clean_client)
      artifact_store = _create_local_artifact_store(clean_client)
  
      with pytest.raises(EntityExistsError):
          clean_client.create_stack(
              name="axels_super_awesome_stack_of_fluffyness",
              components={
                  StackComponentType.ORCHESTRATOR: str(orchestrator.id),
                  StackComponentType.ARTIFACT_STORE: str(artifact_store.id),
              },
          )
  
  
  def test_updating_a_stack_with_new_component_succeeds(clean_client):
      """Tests that updating a new stack with already registered components updates the stack with the new or altered components passed in."""
      stack = _create_local_stack(
          client=clean_client, stack_name="some_new_stack_name"
      )
      clean_client.activate_stack(stack_name_id_or_prefix=stack.name)
  
      old_orchestrator = stack.components[StackComponentType.ORCHESTRATOR][0]
      old_artifact_store = stack.components[StackComponentType.ARTIFACT_STORE][0]
      orchestrator = _create_local_orchestrator(
          client=clean_client, orchestrator_name="different_orchestrator"
      )
  
      with does_not_raise():
          updated_stack = clean_client.update_stack(
              name_id_or_prefix=stack.name,
              component_updates={
                  StackComponentType.ORCHESTRATOR: [str(orchestrator.id)],
              },
          )
  
      active_orchestrator = updated_stack.components[
          StackComponentType.ORCHESTRATOR
      ][0]
      active_artifact_store = updated_stack.components[
          StackComponentType.ARTIFACT_STORE
      ][0]
      assert active_orchestrator != old_orchestrator
      assert active_orchestrator == orchestrator
      assert active_artifact_store == old_artifact_store
  
  
  def test_renaming_stack_with_update_method_succeeds(clean_client):
      """Tests that renaming a stack with the update method succeeds."""
      stack = _create_local_stack(
          client=clean_client, stack_name="some_new_stack_name"
      )
      clean_client.activate_stack(stack.id)
  
      new_stack_name = "new_stack_name"
  
      with does_not_raise():
          clean_client.update_stack(
              name_id_or_prefix=stack.id, name=new_stack_name
          )
      assert clean_client.get_stack(name_id_or_prefix=new_stack_name)
  
  
  def test_register_a_stack_with_unregistered_component_fails(clean_client):
      """Tests that registering a stack with an unregistered component fails."""
      with pytest.raises(KeyError):
          clean_client.create_stack(
              name="axels_empty_stack_of_disappoint",
              components={
                  StackComponentType.ORCHESTRATOR: "orchestrator_doesnt_exist",
                  StackComponentType.ARTIFACT_STORE: "this_also_doesnt",
              },
          )
  
  
  def test_deregistering_the_active_stack(clean_client):
      """Tests that deregistering the active stack fails."""
      with pytest.raises(ValueError):
          clean_client.delete_stack(clean_client.active_stack_model.id)
  
  
  def test_deregistering_a_non_active_stack(clean_client):
      """Tests that deregistering a non-active stack works."""
      stack = _create_local_stack(
          client=clean_client, stack_name="some_new_stack_name"
      )
  
      with does_not_raise():
          clean_client.delete_stack(name_id_or_prefix=stack.id)
  
  
  def test_getting_a_stack_component(clean_client):
      """Tests that getting a stack component returns the correct component."""
      component = clean_client.active_stack_model.components[
          StackComponentType.ORCHESTRATOR
      ][0]
      with does_not_raise():
          registered_component = clean_client.get_stack_component(
              component_type=component.type, name_id_or_prefix=component.id
          )
  
      assert component == registered_component
  
  
  def test_getting_a_nonexisting_stack_component(clean_client):
      """Tests that getting a stack component for a name that isn't registered fails."""
      with pytest.raises(KeyError):
          clean_client.get_stack(name_id_or_prefix=str(uuid4()))
  
  
  def test_registering_a_stack_component_with_existing_name(clean_client):
      """Tests that registering a stack component for an existing name fails."""
      _create_local_orchestrator(
          client=clean_client, orchestrator_name="axels_orchestration_laboratory"
      )
      with pytest.raises(EntityExistsError):
          clean_client.create_stack_component(
              name="axels_orchestration_laboratory",
              flavor="local",
              component_type=StackComponentType.ORCHESTRATOR,
              configuration={},
          )
  
  
  def test_registering_a_new_stack_component_succeeds(clean_client):
      """Tests that registering a stack component works and is persisted."""
      new_artifact_store = _create_local_artifact_store(client=clean_client)
  
      new_client = Client(clean_client.root)
  
      with does_not_raise():
          registered_artifact_store = new_client.get_stack_component(
              component_type=new_artifact_store.type,
              name_id_or_prefix=new_artifact_store.id,
          )
  
      assert registered_artifact_store == new_artifact_store
  
  
  def test_deregistering_a_stack_component_in_stack_fails(clean_client):
      """Tests that deregistering a stack component works and is persisted."""
      component = _create_local_stack(
          clean_client,
          "local_stack",
          orchestrator_name="unregistered_orchestrator",
      ).components[StackComponentType.ORCHESTRATOR][0]
  
      with pytest.raises(IllegalOperationError):
          clean_client.delete_stack_component(
              component_type=StackComponentType.ORCHESTRATOR,
              name_id_or_prefix=str(component.id),
          )
  
  
  def test_deregistering_a_stack_component_that_is_part_of_a_registered_stack(
      clean_client,
  ):
      """Tests that deregistering a stack component that is part of a registered stack fails."""
      component = clean_client.active_stack_model.components[
          StackComponentType.ORCHESTRATOR
      ][0]
  
      with pytest.raises(IllegalOperationError):
          clean_client.delete_stack_component(
              name_id_or_prefix=component.id,
              component_type=StackComponentType.ORCHESTRATOR,
          )
  
  
  def test_getting_a_pipeline(clean_client: "Client"):
      """Tests fetching of a pipeline."""
      # Non-existent ID
      with pytest.raises(KeyError):
          clean_client.get_pipeline(name_id_or_prefix=uuid4())
  
      # Non-existent name
      with pytest.raises(KeyError):
          clean_client.get_pipeline(name_id_or_prefix="non_existent")
  
      request = PipelineRequest(
          project=clean_client.active_project.id,
          name="pipeline",
      )
      response_1 = clean_client.zen_store.create_pipeline(request)
  
      pipeline = clean_client.get_pipeline(name_id_or_prefix=response_1.id)
      assert pipeline == response_1
  
      pipeline = clean_client.get_pipeline(name_id_or_prefix="pipeline")
      assert pipeline == response_1
  
  
  def test_listing_pipelines(clean_client):
      """Tests listing of pipelines."""
      assert clean_client.list_pipelines().total == 0
  
      request = PipelineRequest(
          project=clean_client.active_project.id,
          name="pipeline",
      )
      response_1 = clean_client.zen_store.create_pipeline(request)
      request.name = "other_pipeline"
      clean_client.zen_store.create_pipeline(request)
  
      assert clean_client.list_pipelines().total == 2
  
      assert clean_client.list_pipelines(name="pipeline").total == 1
      assert clean_client.list_pipelines(name="pipeline").items[0] == response_1
  
      assert clean_client.list_pipelines(name="yet_another_pipeline").total == 0
  
  
  def test_create_run_metadata_for_pipeline_run(clean_client_with_run: Client):
      """Test creating run metadata linked only to a pipeline run."""
      pipeline_run = clean_client_with_run.list_pipeline_runs()[0]
      # Assert that the created metadata is correct
      clean_client_with_run.create_run_metadata(
          metadata={"axel": "is awesome"},
          resources=[
              RunMetadataResource(
                  id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
              )
          ],
      )
      rm = clean_client_with_run.get_pipeline_run(pipeline_run.id).run_metadata
  
      assert isinstance(rm, dict)
      assert len(rm.values()) == 1
      assert rm["axel"] == "is awesome"
  
  
  def test_create_run_metadata_for_step_run(clean_client_with_run: Client):
      """Test creating run metadata linked only to a step run."""
      step_run = clean_client_with_run.list_run_steps()[0]
  
      # Assert that the created metadata is correct
      clean_client_with_run.create_run_metadata(
          metadata={"axel": "is awesome"},
          resources=[
              RunMetadataResource(
                  id=step_run.id, type=MetadataResourceTypes.STEP_RUN
              )
          ],
      )
      rm = clean_client_with_run.get_run_step(step_run.id).run_metadata
  
      assert isinstance(rm, dict)
      assert len(rm.values()) == 1
      assert rm["axel"] == "is awesome"
  
  
  def test_create_run_metadata_for_artifact(clean_client_with_run: Client):
      """Test creating run metadata linked to an artifact."""
      artifact_version = clean_client_with_run.list_artifact_versions()[0]
  
      # Assert that the created metadata is correct
      clean_client_with_run.create_run_metadata(
          metadata={"axel": "is awesome"},
          resources=[
              RunMetadataResource(
                  id=artifact_version.id,
                  type=MetadataResourceTypes.ARTIFACT_VERSION,
              )
          ],
      )
  
      rm = clean_client_with_run.get_artifact_version(
          artifact_version.id
      ).run_metadata
  
      assert isinstance(rm, dict)
      assert rm["axel"] == "is awesome"
  
  
  def test_create_run_metadata_overwrite(clean_client_with_run: Client):
      """Test creating run metadata with the same key but different value."""
      pipeline_run = clean_client_with_run.list_pipeline_runs()[0]
  
      clean_client_with_run.create_run_metadata(
          metadata={"axel": "is not awesome"},
          resources=[
              RunMetadataResource(
                  id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
              )
          ],
      )
      clean_client_with_run.create_run_metadata(
          metadata={"axel": "is awesome"},
          resources=[
              RunMetadataResource(
                  id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
              )
          ],
      )
  
      rm = clean_client_with_run.get_pipeline_run(pipeline_run.id).run_metadata
      assert rm["axel"] == "is awesome"
  
  
  def test_create_run_metadata_dict_merge(clean_client_with_run: Client):
      """Test creating run metadata with the same key merges dictionaries."""
      pipeline_run = clean_client_with_run.list_pipeline_runs()[0]
  
      clean_client_with_run.create_run_metadata(
          metadata={"key": {"a": 1}},
          resources=[
              RunMetadataResource(
                  id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
              )
          ],
      )
      clean_client_with_run.create_run_metadata(
          metadata={"key": {"a": 2, "b": 2}},
          resources=[
              RunMetadataResource(
                  id=pipeline_run.id, type=MetadataResourceTypes.PIPELINE_RUN
              )
          ],
      )
  
      rm = clean_client_with_run.get_pipeline_run(pipeline_run.id).run_metadata
      assert rm["key"] == {"a": 2, "b": 2}
  
  
  # .---------.
  # | SECRETS |
  # '---------'
  
  
  def random_secret_name(prefix: str = "aria") -> str:
      """Function to get a random secret name or prefix."""
      return f"pytest_{prefix}_{random_str(4)}"
  
  
  @contextmanager
  def random_secret_context() -> Generator[str, None, None]:
      """Context for testing secrets.
  
      Generates a random secret prefix to avoid conflicts. Yields the prefix.
      After the context is exited, all secrets with that prefix are deleted.
      """
      prefix = random_secret_name()
      yield prefix
  
      client = Client()
      for secret in client.list_secrets(name=f"startswith:{prefix}").items:
          client.delete_secret(secret.id)
  
  
  def test_create_secret_default_scope():
      """Test that secrets are created as public by default."""
      client = Client()
      with random_secret_context() as name:
          s = client.create_secret(
              name=name,
              values={"key": "value"},
          )
          assert s.private is False
          assert s.name == name
          assert s.secret_values == {"key": "value"}
  
  
  def test_create_private_secret():
      """Test creating private secrets."""
      client = Client()
      with random_secret_context() as name:
          s = client.create_secret(
              name=name,
              private=True,
              values={"key": "value"},
          )
          assert s.private is True
          assert s.name == name
          assert s.secret_values == {"key": "value"}
  
  
  def test_create_secret_existing_name_scope():
      """Test that creating a secret with an existing name fails."""
      client = Client()
      with random_secret_context() as name:
          client.create_secret(
              name=name,
              values={"key": "value"},
          )
  
          with pytest.raises(EntityExistsError):
              client.create_secret(
                  name=name,
                  values={"key": "value"},
              )
  
  
  def test_create_private_secret_existing_name():
      """Test that creating a private secret with an existing name fails."""
      client = Client()
      with random_secret_context() as name:
          client.create_secret(
              name=name,
              private=True,
              values={"key": "value"},
          )
  
          with pytest.raises(EntityExistsError):
              client.create_secret(
                  name=name,
                  private=True,
                  values={"key": "value"},
              )
  
  
  def test_create_secret_existing_name_different_scope():
      """Test that creating a secret with the same name in different scopes succeeds."""
      client = Client()
      with random_secret_context() as name:
          s1 = client.create_secret(
              name=name,
              values={"key": "value"},
          )
  
          with does_not_raise():
              s2 = client.create_secret(
                  name=name,
                  private=True,
                  values={"key": "value"},
              )
  
          assert s1.id != s2.id
          assert s1.name == s2.name
  
  
  # ---------------
  # Pipeline Builds
  # ---------------
  
  
  def test_listing_builds(clean_client):
      """Tests listing builds."""
      builds = clean_client.list_builds()
      assert len(builds) == 0
  
      request = PipelineBuildRequest(
          project=clean_client.active_project.id,
          images={},
          is_local=False,
          contains_code=True,
      )
      response = clean_client.zen_store.create_build(request)
  
      builds = clean_client.list_builds()
      assert len(builds) == 1
      assert builds[0] == response
  
      builds = clean_client.list_builds(stack_id=uuid4())
      assert len(builds) == 0
  
  
  def test_getting_builds(clean_client):
      """Tests getting builds."""
  
      with pytest.raises(KeyError):
          clean_client.get_build(str(uuid4()))
  
      request = PipelineBuildRequest(
          project=clean_client.active_project.id,
          images={},
          is_local=False,
          contains_code=True,
      )
      response = clean_client.zen_store.create_build(request)
  
      with does_not_raise():
          build = clean_client.get_build(str(response.id))
  
      assert build == response
  
  
  def test_deleting_builds(clean_client):
      """Tests deleting builds."""
      with pytest.raises(KeyError):
          clean_client.delete_build(str(uuid4()))
  
      request = PipelineBuildRequest(
          project=clean_client.active_project.id,
          images={},
          is_local=False,
          contains_code=True,
      )
      response = clean_client.zen_store.create_build(request)
  
      with does_not_raise():
          clean_client.delete_build(str(response.id))
  
      with pytest.raises(KeyError):
          clean_client.get_build(str(response.id))
  
  
  # --------------------
  # Pipeline Snasphots
  # --------------------
  
  
  def test_listing_snapshots(clean_client):
      """Tests listing snapshots."""
      snapshots = clean_client.list_snapshots(named_only=False)
      assert len(snapshots) == 0
  
      pipeline = clean_client.zen_store.create_pipeline(
          PipelineRequest(
              project=clean_client.active_project.id,
              name="pipeline_name",
          )
      )
  
      request = PipelineSnapshotRequest(
          project=clean_client.active_project.id,
          stack=clean_client.active_stack.id,
          run_name_template="",
          pipeline_configuration={"name": "pipeline_name"},
          client_version="0.12.3",
          server_version="0.12.3",
          pipeline=pipeline.id,
          is_dynamic=False,
      )
      response = clean_client.zen_store.create_snapshot(request)
  
      snapshots = clean_client.list_snapshots(named_only=False)
      assert len(snapshots) == 1
      assert snapshots[0] == response
  
      snapshots = clean_client.list_snapshots(named_only=False, stack=uuid4())
      assert len(snapshots) == 0
  
  
  def test_getting_snapshots(clean_client):
      """Tests getting snapshots."""
      with pytest.raises(KeyError):
          clean_client.get_snapshot(str(uuid4()))
  
      pipeline = clean_client.zen_store.create_pipeline(
          PipelineRequest(
              project=clean_client.active_project.id,
              name="pipeline_name",
          )
      )
      request = PipelineSnapshotRequest(
          project=clean_client.active_project.id,
          stack=clean_client.active_stack.id,
          run_name_template="",
          pipeline_configuration={"name": "pipeline_name"},
          client_version="0.12.3",
          server_version="0.12.3",
          pipeline=pipeline.id,
          is_dynamic=False,
      )
      response = clean_client.zen_store.create_snapshot(request)
  
      with does_not_raise():
          snapshot = clean_client.get_snapshot(str(response.id))
  
      assert snapshot == response
  
  
  def test_deleting_snapshots(clean_client):
      """Tests deleting snapshots."""
      with pytest.raises(KeyError):
          clean_client.delete_snapshot(str(uuid4()))
  
      pipeline = clean_client.zen_store.create_pipeline(
          PipelineRequest(
              project=clean_client.active_project.id,
              name="pipeline_name",
          )
      )
      request = PipelineSnapshotRequest(
          project=clean_client.active_project.id,
          stack=clean_client.active_stack.id,
          run_name_template="",
          pipeline_configuration={"name": "pipeline_name"},
          client_version="0.12.3",
          server_version="0.12.3",
          pipeline=pipeline.id,
          is_dynamic=False,
      )
      response = clean_client.zen_store.create_snapshot(request)
  
      with does_not_raise():
          clean_client.delete_snapshot(str(response.id))
  
      with pytest.raises(KeyError):
          clean_client.get_snapshot(str(response.id))
  
  
  def test_get_run(clean_client: Client, connected_two_step_pipeline):
      """Test that `get_run()` returns the correct run."""
      pipeline_instance = connected_two_step_pipeline(
          step_1=constant_int_output_test_step,
          step_2=int_plus_one_test_step,
      )
      pipeline_instance()
      run_ = clean_client.get_pipeline("connected_two_step_pipeline").runs[0]
      assert clean_client.get_pipeline_run(run_.name) == run_
  
  
  def test_get_run_fails_for_non_existent_run(clean_client: Client):
      """Test that `get_run()` raises a `KeyError` for non-existent runs."""
      with pytest.raises(KeyError):
          clean_client.get_pipeline_run("non_existent_run")
  
  
  class ClientCrudTestConfig(BaseModel):
      entity_name: str
      create_args: Dict[str, Any] = {}
      get_args: Dict[str, Any] = {}
      update_args: Dict[str, Any] = {}
      delete_args: Dict[str, Any] = {}
  
  
  crud_test_configs = [
      ClientCrudTestConfig(
          entity_name="user",
          create_args={"name": sample_name("user_name")},
          update_args={"updated_name": sample_name("updated_user_name")},
      ),
      ClientCrudTestConfig(
          entity_name="project",
          create_args={"name": sample_name("project_name"), "description": ""},
          update_args={"new_name": sample_name("updated_project_name")},
      ),
      ClientCrudTestConfig(
          entity_name="stack",
          create_args={
              "name": sample_name("stack_name"),
              "components": {
                  StackComponentType.ORCHESTRATOR: "default",
                  StackComponentType.ARTIFACT_STORE: "default",
              },
          },
          update_args={"name": sample_name("updated_stack_name")},
      ),
      ClientCrudTestConfig(
          entity_name="stack_component",
          create_args={
              "name": sample_name("stack_component_name"),
              "flavor": "local",
              "component_type": StackComponentType.ORCHESTRATOR,
              "configuration": {},
          },
          get_args={"component_type": StackComponentType.ORCHESTRATOR},
          update_args={
              "name": sample_name("updated_stack_component_name"),
              "component_type": StackComponentType.ORCHESTRATOR,
          },
          delete_args={"component_type": StackComponentType.ORCHESTRATOR},
      ),
      ClientCrudTestConfig(
          entity_name="flavor",
          create_args={
              "source": "tests.unit.test_flavor.AriaOrchestratorFlavor",
              "component_type": StackComponentType.ORCHESTRATOR,
          },
      ),
      ClientCrudTestConfig(
          entity_name="code_repository",
          create_args={
              "name": sample_name("code_repository_name"),
              "config": {},
              "source": Source(
                  module="tests.unit.pipelines.test_build_utils",
                  attribute="StubCodeRepository",
                  type="unknown",
              ),
          },
          update_args={
              "name": sample_name("updated_code_repository_name"),
          },
      ),
  ]
  
  
  @pytest.mark.parametrize(
      "crud_test_config",
      crud_test_configs,
      ids=[c.entity_name for c in crud_test_configs],
  )
  def test_basic_crud_for_entity(
      crud_test_config: ClientCrudTestConfig, clean_client
  ):
      """Tests basic CRUD method on the client."""
      create_method = getattr(
          clean_client, f"create_{crud_test_config.entity_name}"
      )
      get_method = getattr(clean_client, f"get_{crud_test_config.entity_name}")
      delete_method = getattr(
          clean_client, f"delete_{crud_test_config.entity_name}"
      )
  
      entity = create_method(**crud_test_config.create_args)
      try:
          assert hasattr(entity, "name")
          assert len(getattr(entity, "name")) > 1
  
          with does_not_raise():
              get_method(
                  name_id_or_prefix=entity.name, **crud_test_config.get_args
              )
              get_method(
                  name_id_or_prefix=entity.id, **crud_test_config.get_args
              )
              get_method(
                  name_id_or_prefix=entity.name[:-1],
                  allow_name_prefix_match=True,
                  **crud_test_config.get_args,
              )
  
          with pytest.raises(KeyError):
              get_method(
                  name_id_or_prefix=entity.name[:-1],
                  allow_name_prefix_match=False,
                  **crud_test_config.get_args,
              )
  
          if hasattr(clean_client, f"update_{crud_test_config.entity_name}"):
              update_method = getattr(
                  clean_client, f"update_{crud_test_config.entity_name}"
              )
  
              # Updating works with id prefix
              id_prefix = str(entity.id)[:5]
              entity = update_method(
                  name_id_or_prefix=id_prefix, **crud_test_config.update_args
              )
  
              with pytest.raises(KeyError):
                  # Updating doesn't work with name prefix
                  update_method(
                      name_id_or_prefix=entity.name[:-1],
                      **crud_test_config.update_args,
                  )
  
          with pytest.raises(KeyError):
              # Deleting doesn't work with name prefix
              delete_method(
                  name_id_or_prefix=entity.name[:-1],
                  **crud_test_config.delete_args,
              )
  
          # Deleting works with id prefix
          id_prefix = str(entity.id)[:5]
          delete_method(
              name_id_or_prefix=id_prefix, **crud_test_config.delete_args
          )
      finally:
  