text: # Copyright Materialize, Inc. and contributors. All rights reserved.
  #
  # Use of this software is governed by the Business Source License
  # included in the LICENSE file at the root of this repository.
  #
  # As of the Change Date specified in that file, in accordance with
  # the Business Source License, use of this software will be governed
  # by the Apache License, Version 2.0.
  
  import os
  import time
  
  from materialize.mzcompose import Composition
  from materialize.mzcompose.services import (
      Kafka,
      Materialized,
      SchemaRegistry,
      Testdrive,
      Zookeeper,
  )
  
  mz_options = "--persistent-user-tables --persistent-kafka-sources --disable-persistent-system-tables-test"
  
  mz_default = Materialized(options=mz_options)
  
  mz_logical_compaction_window_off = Materialized(
      options=f"{mz_options} --logical-compaction-window=off"
  )
  
  mz_disable_user_indexes = Materialized(
      options=f"{mz_options} --disable-user-indexes",
  )
  
  prerequisites = ["zookeeper", "kafka", "schema-registry"]
  
  SERVICES = [
      Zookeeper(),
      Kafka(),
      SchemaRegistry(),
      mz_default,
      Testdrive(no_reset=True),
  ]
  
  td_test = os.environ.pop("TD_TEST", "*")
  
  
  def workflow_default(c: Composition) -> None:
      # TODO: add back mz_logical_compaction_window_off in the line below.
      # See: https://github.com/MaterializeInc/materialize/issues/10488
      for mz in [mz_default]:
          with c.override(mz):
              workflow_kafka_sources(c)
              workflow_user_tables(c)
              workflow_failpoints(c)
  
      workflow_disable_user_indexes(c)
  
  
  def workflow_kafka_sources(c: Composition) -> None:
      seed = round(time.time())
  
      c.start_and_wait_for_tcp(services=prerequisites)
  
      c.up("materialized")
      c.wait_for_materialized("materialized")
  
      c.run("testdrive-svc", f"--seed={seed}", f"kafka-sources/*{td_test}*-before.td")
  
      c.kill("materialized")
      c.up("materialized")
      c.wait_for_materialized("materialized")
  
      # And restart again, for extra stress
      c.kill("materialized")
      c.up("materialized")
      c.wait_for_materialized("materialized")
  
      c.run("testdrive-svc", f"--seed={seed}", f"kafka-sources/*{td_test}*-after.td")
  
      # Do one more restart, just in case and just confirm that Mz is able to come up
      c.kill("materialized")
      c.up("materialized")
      c.wait_for_materialized("materialized")
  
      c.kill("materialized")
      c.rm("materialized", "testdrive-svc", destroy_volumes=True)
      c.rm_volumes("mzdata")
  
  
  def workflow_user_tables(c: Composition) -> None:
      seed = round(time.time())
  
      c.up("materialized")
      c.wait_for_materialized()
  
      c.run(
          "testdrive-svc",
          f"--seed={seed}",
          f"user-tables/table-persistence-before-{td_test}.td",
      )
  
      c.kill("materialized")
      c.up("materialized")
  
      c.run(
          "testdrive-svc",
          f"--seed={seed}",
          f"user-tables/table-persistence-after-{td_test}.td",
      )
  
      c.kill("materialized")
      c.rm("materialized", "testdrive-svc", destroy_volumes=True)
      c.rm_volumes("mzdata")
  
  
  def workflow_failpoints(c: Composition) -> None:
      seed = round(time.time())
  
      c.up("materialized")
      c.wait_for_materialized()
  
      c.run("testdrive-svc", f"--seed={seed}", f"failpoints/{td_test}.td")
  
      c.kill("materialized")
      c.rm("materialized", "testdrive-svc", destroy_volumes=True)
      c.rm_volumes("mzdata")
  
  
  def workflow_disable_user_indexes(c: Composition) -> None:
      seed = round(time.time())
  
      c.start_and_wait_for_tcp(services=prerequisites)
  
      c.up("materialized")
      c.wait_for_materialized()
  
      c.run("testdrive-svc", f"--seed={seed}", "disable-user-indexes/before.td")
  
      c.kill("materialized")
  
      with c.override(mz_disable_user_indexes):
          c.up("materialized")
          c.wait_for_materialized()
  
          c.run("testdrive-svc", f"--seed={seed}", "disable-user-indexes/after.td")
  
          c.kill("materialized")
  
      c.rm("materialized", "testdrive-svc", destroy_volumes=True)
  
      c.rm_volumes("mzdata")
  