Transforms and Extensions

/**
 * Trait providing test transformation capabilities (mixed into BaseFunSuite)
 */
trait TestTransforms {
  /**
   * Get the list of test transforms to apply
   * Override this method to customize test transformation
   */
  def munitTestTransforms: List[TestTransform]
  
  /**
   * Apply test transforms to a single test
   * This method is called automatically for each test
   */
  def munitTestTransform(test: Test): Test
  
  /** Check if flaky tests should be treated as passing */
  def munitFlakyOK: Boolean
}

/**
 * A transformation that can be applied to individual tests
 * @param name Descriptive name for the transform
 * @param fn Function that transforms a Test into a modified Test
 */
final class TestTransform(val name: String, fn: Test => Test) extends Function1[Test, Test] {
  def apply(test: Test): Test = fn(test)
}

/** Transform that handles tests marked with the Fail tag */
def munitFailTransform: TestTransform

/** Transform that handles tests marked with the Flaky tag */
def munitFlakyTransform: TestTransform

/**
 * Transform that appends additional context to failure messages
 * @param buildSuffix Function to generate additional context for a test
 */
def munitAppendToFailureMessage(buildSuffix: Test => Option[String]): TestTransform

class TransformExamples extends FunSuite {
  // Custom transform to add timing information
  val timingTransform = new TestTransform(
    "timing",
    test => test.withBodyMap { originalBody =>
      val start = System.currentTimeMillis()
      originalBody.andThen { _ =>
        val duration = System.currentTimeMillis() - start
        println(s"Test '${test.name}' took ${duration}ms")
        Future.unit
      }
    }
  )
  
  // Custom transform to retry flaky tests
  val retryTransform = new TestTransform(
    "retry",
    test => {
      if (test.tags.contains(new Tag("retry"))) {
        test.withBodyMap { originalBody =>
          originalBody.recoverWith {
            case _: AssertionError =>
              println(s"Retrying test: ${test.name}")
              originalBody
          }
        }
      } else {
        test
      }
    }
  )
  
  override def munitTestTransforms: List[TestTransform] = 
    super.munitTestTransforms ++ List(timingTransform, retryTransform)
  
  test("normal test") {
    assertEquals(1 + 1, 2)
  }
  
  test("flaky network test".tag(new Tag("retry"))) {
    // This test will be retried once if it fails
    callExternalAPI()
  }
}

/**
 * Trait providing suite transformation capabilities (mixed into BaseFunSuite)
 */
trait SuiteTransforms {
  /**
   * Get the list of suite transforms to apply
   * Override this method to customize suite transformation
   */
  def munitSuiteTransforms: List[SuiteTransform]
  
  /**
   * Apply suite transforms to the entire test list
   * This method is called automatically when collecting tests
   */
  def munitSuiteTransform(tests: List[Test]): List[Test]
  
  /** Override to ignore the entire test suite */
  def munitIgnore: Boolean = false
  
  /** Check if running in continuous integration environment */
  def isCI: Boolean
}

/**
 * A transformation that can be applied to entire test suites
 * @param name Descriptive name for the transform
 * @param fn Function that transforms a List[Test] into a modified List[Test]
 */
final class SuiteTransform(val name: String, fn: List[Test] => List[Test]) extends Function1[List[Test], List[Test]] {
  def apply(tests: List[Test]): List[Test] = fn(tests)
}

/** Transform that handles ignored test suites */
def munitIgnoreSuiteTransform: SuiteTransform

/** Transform that handles the Only tag (runs only marked tests) */
def munitOnlySuiteTransform: SuiteTransform

class SuiteTransformExamples extends FunSuite {
  // Custom transform to randomize test order
  val randomizeTransform = new SuiteTransform(
    "randomize",
    tests => scala.util.Random.shuffle(tests)
  )
  
  // Custom transform to group tests by tags
  val groupByTagTransform = new SuiteTransform(
    "groupByTag",
    tests => {
      val (slowTests, fastTests) = tests.partition(_.tags.contains(Slow))
      fastTests ++ slowTests // Run fast tests first
    }
  )
  
  // Custom transform to skip integration tests in CI
  val skipIntegrationInCI = new SuiteTransform(
    "skipIntegrationInCI",
    tests => {
      if (isCI) {
        tests.filterNot(_.tags.contains(new Tag("integration")))
      } else {
        tests
      }
    }
  )
  
  override def munitSuiteTransforms: List[SuiteTransform] = 
    super.munitSuiteTransforms ++ List(
      groupByTagTransform,
      skipIntegrationInCI
    )
  
  test("fast unit test") {
    assertEquals(1 + 1, 2)
  }
  
  test("slow integration test".tag(Slow).tag(new Tag("integration"))) {
    // This test may be skipped in CI
    performIntegrationTest()
  }
}

/**
 * Trait providing value transformation capabilities (mixed into BaseFunSuite)
 */
trait ValueTransforms {
  /**
   * Get the list of value transforms to apply
   * Override this method to customize value transformation
   */
  def munitValueTransforms: List[ValueTransform]
  
  /**
   * Apply value transforms to convert any test return value to Future[Any]
   * This method is called automatically for each test body
   */
  def munitValueTransform(testValue: => Any): Future[Any]
}

/**
 * A transformation that can be applied to test return values
 * @param name Descriptive name for the transform
 * @param fn Partial function that transforms specific types to Future[Any]
 */
final class ValueTransform(val name: String, fn: PartialFunction[Any, Future[Any]]) extends Function1[Any, Option[Future[Any]]] {
  def apply(value: Any): Option[Future[Any]] = fn.lift(value)
}

/** Transform that handles Future return values */
def munitFutureTransform: ValueTransform

import scala.util.{Try, Success, Failure}
import scala.concurrent.Future

class ValueTransformExamples extends FunSuite {
  // Custom transform to handle Try values
  val tryTransform = new ValueTransform(
    "try",
    {
      case Success(value) => Future.successful(value)
      case Failure(exception) => Future.failed(exception)
    }
  )
  
  // Custom transform to handle custom Result type
  case class Result[T](value: T, errors: List[String]) {
    def isSuccess: Boolean = errors.isEmpty
  }
  
  val resultTransform = new ValueTransform(
    "result",
    {
      case result: Result[_] =>
        if (result.isSuccess) {
          Future.successful(result.value)
        } else {
          Future.failed(new AssertionError(s"Result failed: ${result.errors.mkString(", ")}"))
        }
    }
  )
  
  override def munitValueTransforms: List[ValueTransform] =
    super.munitValueTransforms ++ List(tryTransform, resultTransform)
  
  test("test returning Try") {
    Try {
      assertEquals(2 + 2, 4)
      "success"
    }
  }
  
  test("test returning custom Result") {
    val data = processData()
    Result(data, List.empty) // Success case
  }
  
  test("test returning failed Result") {
    Result("", List("validation error")) // This will fail the test
  }
}

class ConditionalTransforms extends FunSuite {
  // Transform that only applies to tests with specific tags
  val debugTransform = new TestTransform(
    "debug",
    test => {
      if (test.tags.contains(new Tag("debug"))) {
        test.withBodyMap { originalBody =>
          println(s"[DEBUG] Starting test: ${test.name}")
          originalBody.andThen { result =>
            println(s"[DEBUG] Completed test: ${test.name}")
            result
          }
        }
      } else {
        test
      }
    }
  )
  
  // Transform that modifies behavior based on environment
  val environmentTransform = new SuiteTransform(
    "environment",
    tests => {
      val environment = System.getProperty("test.environment", "local")
      environment match {
        case "ci" => tests.filterNot(_.tags.contains(new Tag("local-only")))
        case "staging" => tests.filterNot(_.tags.contains(new Tag("prod-only")))
        case _ => tests
      }
    }
  )
}

class MetricsTransforms extends FunSuite {
  private val testMetrics = scala.collection.mutable.Map[String, Long]()
  
  val metricsTransform = new TestTransform(
    "metrics",
    test => test.withBodyMap { originalBody =>
      val start = System.nanoTime()
      originalBody.andThen { result =>
        val duration = System.nanoTime() - start
        testMetrics(test.name) = duration
        result
      }
    }
  )
  
  override def munitTestTransforms = List(metricsTransform)
  
  override def afterAll(): Unit = {
    super.afterAll()
    println("Test Execution Times:")
    testMetrics.toList.sortBy(_._2).foreach { case (name, nanos) =>
      println(f"  $name: ${nanos / 1000000.0}%.2f ms")
    }
  }
}

class ErrorContextTransforms extends FunSuite {
  val contextTransform = new TestTransform(
    "context",
    test => test.withBodyMap { originalBody =>
      originalBody.recoverWith {
        case e: AssertionError =>
          val enhancedMessage = s"${e.getMessage}\n" +
            s"Test: ${test.name}\n" +
            s"Tags: ${test.tags.map(_.value).mkString(", ")}\n" +
            s"Location: ${test.location}"
          Future.failed(new AssertionError(enhancedMessage, e))
      }
    }
  )
  
  override def munitTestTransforms = List(contextTransform)
}

class EnvironmentBasedSuite extends FunSuite {
  override def munitFlakyOK: Boolean = {
    // Allow flaky tests in development but not in CI
    !isCI && System.getProperty("munit.flaky.ok", "false").toBoolean
  }
  
  override def isCI: Boolean = {
    System.getenv("CI") != null || 
    System.getProperty("CI") != null
  }
}

override def munitTestTransforms: List[TestTransform] = List(
  timingTransform,      // Applied first
  retryTransform,       // Applied second  
  loggingTransform      // Applied last
) ++ super.munitTestTransforms // Include built-in transforms