TestHelper.kt

package io.sentry.systemtest.util

import com.apollographql.apollo.api.ApolloResponse
import com.apollographql.apollo.api.Operation
import io.sentry.JsonSerializer
import io.sentry.ProfileChunk
import io.sentry.SentryEnvelopeHeader
import io.sentry.SentryEvent
import io.sentry.SentryItemType
import io.sentry.SentryLogEvents
import io.sentry.SentryMetricsEvents
import io.sentry.SentryOptions
import io.sentry.protocol.FeatureFlag
import io.sentry.protocol.SentrySpan
import io.sentry.protocol.SentryTransaction
import io.sentry.systemtest.graphql.GraphqlTestClient
import java.io.BufferedReader
import java.io.ByteArrayInputStream
import java.io.File
import java.io.InputStreamReader
import java.io.PrintWriter
import kotlin.test.assertEquals
import kotlin.test.assertFalse
import kotlin.test.assertNotNull
import kotlin.test.assertTrue

class TestHelper(backendUrl: String) {
  val restClient: RestTestClient
  val graphqlClient: GraphqlTestClient
  val sentryClient: SentryMockServerClient
  val jsonSerializer: JsonSerializer
  val dsn = "http://502f25099c204a2fbf4cb16edc5975d1@localhost:8000/0"

  var envelopeCounts: EnvelopeCounts? = null

  init {
    restClient = RestTestClient(backendUrl)
    sentryClient = SentryMockServerClient("http://localhost:8000")
    graphqlClient = GraphqlTestClient(backendUrl)
    jsonSerializer = JsonSerializer(SentryOptions.empty())
  }

  fun snapshotEnvelopeCount() {
    envelopeCounts = sentryClient.getEnvelopeCount()
  }

  fun ensureEnvelopeCountIncreased() {
    Thread.sleep(1000)
    val envelopeCountsAfter = sentryClient.getEnvelopeCount()
    assertTrue(envelopeCountsAfter!!.envelopes!! > envelopeCounts!!.envelopes!!)
  }

  fun ensureEnvelopeReceived(retryCount: Int = 1, callback: ((String) -> Boolean)) {
    val envelopes = sentryClient.getEnvelopes()
    assertNotNull(envelopes.envelopes)
    envelopes.envelopes.forEach { envelopeString ->
      val didMatch = callback(envelopeString)
      if (didMatch) {
        return
      }
    }
    if (retryCount <= 0) {
      throw RuntimeException("Unable to find matching envelope received by relay")
    } else {
      Thread.sleep(10000)
      ensureEnvelopeReceived(retryCount - 1, callback)
    }
  }

  fun ensureNoEnvelopeReceived(callback: ((String) -> Boolean)) {
    Thread.sleep(10000)
    val envelopes = sentryClient.getEnvelopes()

    if (envelopes.envelopes.isNullOrEmpty()) {
      return
    }

    envelopes.envelopes.forEach { envelopeString ->
      val didMatch = callback(envelopeString)
      if (didMatch) {
        throw RuntimeException("Found unexpected matching envelope received by relay")
      }
    }
  }

  fun ensureTransactionReceived(callback: ((SentryTransaction, SentryEnvelopeHeader) -> Boolean)) {
    ensureEnvelopeReceived { envelopeString -> checkIfTransactionMatches(envelopeString, callback) }
  }

  fun ensureProfileChunkReceived(callback: ((ProfileChunk, SentryEnvelopeHeader) -> Boolean)) {
    ensureEnvelopeReceived { envelopeString -> checkIfProfileMatches(envelopeString, callback) }
  }

  fun ensureNoTransactionReceived(
    callback: ((SentryTransaction, SentryEnvelopeHeader) -> Boolean)
  ) {
    ensureNoEnvelopeReceived { envelopeString ->
      checkIfTransactionMatches(envelopeString, callback)
    }
  }

  fun ensureLogsReceived(callback: ((SentryLogEvents, SentryEnvelopeHeader) -> Boolean)) {
    ensureEnvelopeReceived { envelopeString -> checkIfLogsMatch(envelopeString, callback) }
  }

  private fun checkIfLogsMatch(
    envelopeString: String,
    callback: ((SentryLogEvents, SentryEnvelopeHeader) -> Boolean),
  ): Boolean {
    val deserializeEnvelope = jsonSerializer.deserializeEnvelope(envelopeString.byteInputStream())
    if (deserializeEnvelope == null) {
      return false
    }

    val envelopeHeader = deserializeEnvelope.header

    val logsItem = deserializeEnvelope.items.firstOrNull { it.header.type == SentryItemType.Log }
    if (logsItem == null) {
      return false
    }

    val logs = logsItem.getLogs(jsonSerializer)
    if (logs == null) {
      return false
    }

    return callback(logs, envelopeHeader)
  }

  fun doesContainMetric(
    metrics: SentryMetricsEvents,
    name: String,
    type: String,
    value: Double,
    unit: String? = null,
  ): Boolean {
    val metricItem =
      metrics.items.firstOrNull { metricItem ->
        metricItem.name == name &&
          metricItem.type == type &&
          metricItem.value == value &&
          (unit == null || metricItem.unit == unit)
      }
    if (metricItem == null) {
      println("Unable to find metric item with name $name, type $type and value $value in metrics:")
      logObject(metrics)
      return false
    }

    return true
  }

  fun ensureMetricsReceived(callback: ((SentryMetricsEvents, SentryEnvelopeHeader) -> Boolean)) {
    ensureEnvelopeReceived { envelopeString -> checkIfMetricsMatch(envelopeString, callback) }
  }

  private fun checkIfMetricsMatch(
    envelopeString: String,
    callback: ((SentryMetricsEvents, SentryEnvelopeHeader) -> Boolean),
  ): Boolean {
    val deserializeEnvelope = jsonSerializer.deserializeEnvelope(envelopeString.byteInputStream())
    if (deserializeEnvelope == null) {
      return false
    }

    val envelopeHeader = deserializeEnvelope.header

    val metricsItem =
      deserializeEnvelope.items.firstOrNull { it.header.type == SentryItemType.TraceMetric }
    if (metricsItem == null) {
      return false
    }

    val metrics = metricsItem.getMetrics(jsonSerializer)
    if (metrics == null) {
      return false
    }

    return callback(metrics, envelopeHeader)
  }

  fun doesContainLogWithBody(logs: SentryLogEvents, body: String): Boolean {
    val logItem = logs.items.firstOrNull { logItem -> logItem.body == body }
    if (logItem == null) {
      println("Unable to find log item with body $body in logs:")
      logObject(logs)
      return false
    }

    return true
  }

  fun doesLogWithBodyHaveAttribute(
    logs: SentryLogEvents,
    body: String,
    attributeKey: String,
    attributeValue: Any?,
  ): Boolean {
    val logItem = logs.items.firstOrNull { logItem -> logItem.body == body }
    if (logItem == null) {
      println("Unable to find log item with body $body in logs:")
      logObject(logs)
      return false
    }

    val attr = logItem.attributes?.get(attributeKey)
    if (attr == null) {
      println("Unable to find attribute $attributeKey on log with body $body:")
      logObject(logItem)
      return false
    }

    if (attr.value != attributeValue) {
      println(
        "Attribute $attributeKey has value ${attr.value} but expected $attributeValue on log with body $body:"
      )
      logObject(logItem)
      return false
    }

    return true
  }

  fun doesMetricHaveAttribute(
    metrics: SentryMetricsEvents,
    metricName: String,
    attributeKey: String,
    attributeValue: Any?,
  ): Boolean {
    val metricItem = metrics.items.firstOrNull { it.name == metricName }
    if (metricItem == null) {
      println("Unable to find metric with name $metricName in metrics:")
      logObject(metrics)
      return false
    }

    val attr = metricItem.attributes?.get(attributeKey)
    if (attr == null) {
      println("Unable to find attribute $attributeKey on metric $metricName:")
      logObject(metricItem)
      return false
    }

    if (attr.value != attributeValue) {
      println(
        "Attribute $attributeKey has value ${attr.value} but expected $attributeValue on metric $metricName:"
      )
      logObject(metricItem)
      return false
    }

    return true
  }

  private fun checkIfTransactionMatches(
    envelopeString: String,
    callback: ((SentryTransaction, SentryEnvelopeHeader) -> Boolean),
  ): Boolean {
    val deserializeEnvelope = jsonSerializer.deserializeEnvelope(envelopeString.byteInputStream())
    if (deserializeEnvelope == null) {
      return false
    }

    val envelopeHeader = deserializeEnvelope.header

    val transactionItem =
      deserializeEnvelope.items.firstOrNull { it.header.type == SentryItemType.Transaction }
    if (transactionItem == null) {
      return false
    }

    val transaction = transactionItem.getTransaction(jsonSerializer)
    if (transaction == null) {
      return false
    }

    return callback(transaction, envelopeHeader)
  }

  private fun checkIfProfileMatches(
    envelopeString: String,
    callback: ((ProfileChunk, SentryEnvelopeHeader) -> Boolean),
  ): Boolean {
    val deserializeEnvelope = jsonSerializer.deserializeEnvelope(envelopeString.byteInputStream())
    if (deserializeEnvelope == null) {
      return false
    }

    val envelopeHeader = deserializeEnvelope.header

    val profileChunkItem =
      deserializeEnvelope.items.firstOrNull { it.header.type == SentryItemType.ProfileChunk }

    if (profileChunkItem == null) {
      return false
    }

    val chunk =
      BufferedReader(InputStreamReader(ByteArrayInputStream(profileChunkItem.data), Charsets.UTF_8))
        .use { eventReader -> jsonSerializer.deserialize(eventReader, ProfileChunk::class.java) }

    if (chunk == null) {
      return false
    }

    return callback(chunk, envelopeHeader)
  }

  fun ensureErrorReceived(callback: ((SentryEvent) -> Boolean)) {
    ensureEnvelopeReceived(retryCount = 3) { envelopeString ->
      val deserializeEnvelope = jsonSerializer.deserializeEnvelope(envelopeString.byteInputStream())
      if (deserializeEnvelope == null) {
        return@ensureEnvelopeReceived false
      }

      val errorItem =
        deserializeEnvelope.items.firstOrNull { it.header.type == SentryItemType.Event }
      if (errorItem == null) {
        return@ensureEnvelopeReceived false
      }

      val error = errorItem.getEvent(jsonSerializer)
      if (error == null) {
        return@ensureEnvelopeReceived false
      }

      val callbackResult = callback(error)
      if (!callbackResult) {
        println("found an error event but it did not match:")
        logObject(error)
      }
      callbackResult
    }
  }

  fun ensureTransactionWithSpanReceived(callback: ((SentrySpan) -> Boolean)) {
    ensureTransactionReceived { transaction, envelopeHeader ->
      transaction.spans.forEach { span ->
        val callbackResult = callback(span)
        if (callbackResult) {
          return@ensureTransactionReceived true
        }
      }
      false
    }
  }

  fun reset() {
    sentryClient.reset()
  }

  fun logObject(obj: Any?) {
    obj ?: return
    PrintWriter(System.out).use { jsonSerializer.serialize(obj, it) }
    println()
  }

  fun <T : Operation.Data> ensureNoErrors(response: ApolloResponse<T>?) {
    response ?: throw RuntimeException("no response")
    assertFalse(response.hasErrors())
  }

  fun <T : Operation.Data> ensureErrorCount(response: ApolloResponse<T>?, errorCount: Int) {
    response ?: throw RuntimeException("no response")
    assertEquals(errorCount, response.errors?.size)
  }

  fun doesTransactionContainSpanWithOp(transaction: SentryTransaction, op: String): Boolean {
    val span = transaction.spans.firstOrNull { span -> span.op == op }
    if (span == null) {
      println("Unable to find span with op $op in transaction:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesTransactionContainSpanWithOpAndDescription(
    transaction: SentryTransaction,
    op: String,
    description: String,
  ): Boolean {
    val span =
      transaction.spans.firstOrNull { span -> span.op == op && span.description == description }
    if (span == null) {
      println("Unable to find span with op $op and description $description in transaction:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesTransactionContainSpanWithDescription(
    transaction: SentryTransaction,
    description: String,
  ): Boolean {
    val span = transaction.spans.firstOrNull { span -> span.description == description }
    if (span == null) {
      println("Unable to find span with description $description in transaction:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesTransactionHaveTraceId(transaction: SentryTransaction, traceId: String): Boolean {
    val spanContext = transaction.contexts.trace
    if (spanContext?.traceId?.toString() != traceId) {
      println("Unable to find trace ID $traceId in transaction:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesTransactionHaveOp(transaction: SentryTransaction, op: String): Boolean {
    val matches = transaction.contexts.trace?.operation == op
    if (!matches) {
      println("Unable to find transaction with op $op:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesTransactionHave(
    transaction: SentryTransaction,
    op: String,
    featureFlag: FeatureFlag? = null,
  ): Boolean {
    val matches = transaction.contexts.trace?.operation == op
    if (!matches) {
      println("Unable to find transaction with op $op:")
      logObject(transaction)
      return false
    }

    val foundFlag = transaction.contexts.trace?.data?.get(featureFlag?.flag)
    if (featureFlag != null && foundFlag == null) {
      println("Unable to find span with feature flag ${featureFlag?.flag}:")
      logObject(transaction)
      return false
    }
    if (featureFlag != null && foundFlag != featureFlag.result) {
      println("Feature flag ${featureFlag?.flag} has unexpected result ${foundFlag}:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesTransactionHaveSpanWith(
    transaction: SentryTransaction,
    op: String,
    featureFlag: FeatureFlag? = null,
    noFeatureFlags: Boolean = false,
  ): Boolean {
    val foundSpan = transaction.spans.firstOrNull { span -> span.op == op }
    if (foundSpan == null) {
      println("Unable to find span with op $op:")
      logObject(transaction)
      return false
    }

    val featureFlagNames =
      foundSpan.data?.keys?.filter { it.startsWith("flag.evaluation.") } ?: emptyList()
    if (noFeatureFlags && featureFlagNames.isNotEmpty()) {
      println("Expected 0 feature flags but found ${featureFlagNames}:")
      logObject(transaction)
      return false
    }

    val foundFlag = foundSpan.data?.get(featureFlag?.flag)
    if (featureFlag != null && foundFlag == null) {
      println("Unable to find span with feature flag ${featureFlag?.flag}:")
      logObject(transaction)
      return false
    }
    if (featureFlag != null && foundFlag != featureFlag.result) {
      println("Feature flag ${featureFlag?.flag} has unexpected result ${foundFlag}:")
      logObject(transaction)
      return false
    }

    return true
  }

  fun doesEventHaveExceptionMessage(event: SentryEvent, expectedMessage: String): Boolean {
    val exceptions = event.exceptions
    if (exceptions == null) {
      println("Unable to find exceptions in event")
      return false
    }

    val foundException = exceptions.firstOrNull { expectedMessage == it.value }
    return foundException != null
  }

  fun doesEventHaveFlag(event: SentryEvent, flag: String, result: Boolean): Boolean {
    val featureFlags = event.contexts.featureFlags
    if (featureFlags == null) {
      println("Unable to find feature flags in event:")
      return false
    }
    val foundFlag =
      featureFlags.values.firstOrNull { featureFlag ->
        println("checking flag ${featureFlag.flag}:${featureFlag.result}")
        featureFlag.flag == flag && featureFlag.result == result
      }
    return foundFlag != null
  }

  fun findJar(prefix: String, inDir: String = "build/libs"): File {
    val buildDir = File(inDir)
    val jarFiles =
      buildDir.listFiles { _, name -> name.startsWith(prefix) && name.endsWith(".jar") }?.toList()
        ?: emptyList()

    if (jarFiles.isEmpty()) {
      throw AssertionError("No JAR found in ${buildDir.absolutePath}")
    }

    return jarFiles.maxOf { it }
  }

  fun launch(jar: File, env: Map<String, String>, enableOtelAutoConfig: Boolean = false): Process {
    val processBuilderList = mutableListOf("java", "--add-opens", "java.base/java.lang=ALL-UNNAMED")

    if (enableOtelAutoConfig) {
      processBuilderList.add("-Dotel.java.global-autoconfigure.enabled=true")
    }

    processBuilderList.add("-jar")
    processBuilderList.add(jar.absolutePath)

    val processBuilder =
      ProcessBuilder(processBuilderList).inheritIO() // forward i/o to current process

    processBuilder.environment().putAll(env)

    return processBuilder.start()
  }
}