From cf912c661302f1a77f7da886846f24fe740564b7 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Mon, 3 Nov 2025 10:08:31 +0100 Subject: [PATCH 1/9] [NA] [BE] Upgrade MySQL container from Testcontainers --- .../comet/opik/api/resources/utils/MigrationUtils.java | 4 ++-- .../opik/api/resources/utils/MySQLContainerUtils.java | 8 ++++---- .../api/resources/v1/events/DatasetEventListenerTest.java | 4 ++-- .../api/resources/v1/events/OnlineScoringEngineTest.java | 4 ++-- .../resources/v1/events/WebhookSubscriberLoggingTest.java | 4 ++-- .../opik/api/resources/v1/internal/UsageResourceTest.java | 4 ++-- .../api/resources/v1/jobs/TraceThreadsClosingJobTest.java | 4 ++-- .../opik/api/resources/v1/priv/AlertResourceTest.java | 4 ++-- .../resources/v1/priv/AnnotationQueuesResourceTest.java | 4 ++-- .../resources/v1/priv/AttachmentResourceMinIOTest.java | 4 ++-- .../api/resources/v1/priv/AttachmentResourceTest.java | 4 ++-- .../api/resources/v1/priv/AuthenticationResourceTest.java | 4 ++-- .../v1/priv/AutomationRuleEvaluatorsResourceTest.java | 4 ++-- .../resources/v1/priv/ChatCompletionsResourceTest.java | 4 ++-- .../api/resources/v1/priv/DatasetExperimentE2ETest.java | 4 ++-- .../opik/api/resources/v1/priv/DatasetsResourceTest.java | 4 ++-- .../api/resources/v1/priv/ExperimentsResourceTest.java | 4 ++-- .../resources/v1/priv/FeedbackDefinitionResourceTest.java | 4 ++-- .../api/resources/v1/priv/GuardrailsResourceTest.java | 4 ++-- .../resources/v1/priv/LlmProviderApiKeyResourceTest.java | 4 ++-- .../resources/v1/priv/ManualEvaluationResourceTest.java | 4 ++-- .../v1/priv/MultiValueFeedbackScoresE2ETest.java | 4 ++-- .../api/resources/v1/priv/OpenTelemetryResourceTest.java | 4 ++-- .../api/resources/v1/priv/OptimizationsResourceTest.java | 4 ++-- .../api/resources/v1/priv/ProjectMetricsResourceTest.java | 4 ++-- .../opik/api/resources/v1/priv/ProjectsResourceTest.java | 4 ++-- .../opik/api/resources/v1/priv/PromptResourceTest.java | 4 ++-- .../opik/api/resources/v1/priv/SpansResourceTest.java | 4 ++-- .../opik/api/resources/v1/priv/TracesResourceTest.java | 4 ++-- .../api/resources/v1/priv/WorkspacesResourceTest.java | 4 ++-- .../api/resources/v1/session/RedirectResourceTest.java | 4 ++-- .../opik/infrastructure/HealthCheckIntegrationTest.java | 4 ++-- .../opik/infrastructure/auth/AuthModuleCache2E2Test.java | 4 ++-- .../auth/AuthModuleNoAuthIntegrationTest.java | 4 ++-- .../comet/opik/infrastructure/bi/BiEventListenerTest.java | 4 ++-- .../opik/infrastructure/bi/DailyUsageReportJobTest.java | 6 +++--- .../bi/OpikGuiceyLifecycleEventListenerTest.java | 4 ++-- .../comet/opik/infrastructure/cache/CacheManagerTest.java | 4 ++-- .../comet/opik/infrastructure/health/IsAliveE2ETest.java | 4 ++-- .../infrastructure/http/cors/CorsDisabledE2ETest.java | 4 ++-- .../opik/infrastructure/http/cors/CorsEnabledE2ETest.java | 4 ++-- .../opik/infrastructure/ratelimit/RateLimitE2ETest.java | 4 ++-- .../redis/RedissonLockServiceIntegrationTest.java | 4 ++-- 43 files changed, 89 insertions(+), 89 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MigrationUtils.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MigrationUtils.java index 20aa5fd14ba..41743a1c31c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MigrationUtils.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MigrationUtils.java @@ -10,7 +10,7 @@ import lombok.experimental.UtilityClass; import org.jdbi.v3.core.Jdbi; import org.testcontainers.clickhouse.ClickHouseContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import ru.yandex.clickhouse.ClickHouseConnectionImpl; import java.sql.SQLException; @@ -29,7 +29,7 @@ public static void runMysqlDbMigration(Jdbi jdbi) { } } - public static void runMysqlDbMigration(MySQLContainer mysqlContainer) { + public static void runMysqlDbMigration(MySQLContainer mysqlContainer) { try (var connection = mysqlContainer.createConnection("")) { runDbMigration(MYSQL_CHANGELOG_FILE, MySQLContainerUtils.migrationParameters(), new JdbcConnection(connection)); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MySQLContainerUtils.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MySQLContainerUtils.java index 4d9183136e3..34cbd1a73b8 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MySQLContainerUtils.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MySQLContainerUtils.java @@ -1,18 +1,18 @@ package com.comet.opik.api.resources.utils; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.utility.DockerImageName; import java.util.Map; public class MySQLContainerUtils { - public static MySQLContainer newMySQLContainer() { + public static MySQLContainer newMySQLContainer() { return newMySQLContainer(true); } - public static MySQLContainer newMySQLContainer(boolean reusable) { - return new MySQLContainer<>(DockerImageName.parse("mysql:8.4.2")) + public static MySQLContainer newMySQLContainer(boolean reusable) { + return new MySQLContainer(DockerImageName.parse("mysql:8.4.2")) .withUrlParam("createDatabaseIfNotExist", "true") .withUrlParam("rewriteBatchedStatements", "true") .withUrlParam("serverTimezone", "UTC") diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java index 274f637c7dc..c53a81d8d61 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java @@ -33,7 +33,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -62,7 +62,7 @@ class DatasetEventListenerTest { private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java index b023e3f01ee..9ce1f9d43ef 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java @@ -70,7 +70,7 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; @@ -223,7 +223,7 @@ class OnlineScoringEngineTest { .formatted(EDGE_CASE_TEMPLATE).trim(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/WebhookSubscriberLoggingTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/WebhookSubscriberLoggingTest.java index df17a8be7f5..fb8d136cf38 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/WebhookSubscriberLoggingTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/WebhookSubscriberLoggingTest.java @@ -37,8 +37,8 @@ import org.redisson.api.RedissonReactiveClient; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.test.StepVerifier; import reactor.util.context.Context; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -71,7 +71,7 @@ class WebhookSubscriberLoggingTest { private static final int AWAIT_POLL_INTERVAL_MS = 500; private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java index 0455d7b109f..a8b3882ce7d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java @@ -35,8 +35,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -71,7 +71,7 @@ class UsageResourceTest { private final String USER = UUID.randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java index dba214898b0..d84f9ccf47d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java @@ -35,7 +35,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; @@ -67,7 +67,7 @@ class TraceThreadsClosingJobTest { private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java index c397be952c5..b2a6c75d06f 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java @@ -73,7 +73,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -144,7 +144,7 @@ class AlertResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); @RegisterApp private final TestDropwizardAppExtension APP; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java index af20aeb0d1a..26029c1f022 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java @@ -47,7 +47,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; @@ -92,7 +92,7 @@ class AnnotationQueuesResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java index 179670888f5..02c903976fd 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java @@ -44,7 +44,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -79,7 +79,7 @@ class AttachmentResourceMinIOTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); private final GenericContainer ZOOKEEPER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer MINIO = MinIOContainerUtils.newMinIOContainer(); @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java index 147c1492139..8f78ca76da0 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java @@ -36,7 +36,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -66,7 +66,7 @@ class AttachmentResourceTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); private final GenericContainer ZOOKEEPER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer MINIO = MinIOContainerUtils.newMinIOContainer(); @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java index f657a1518ec..107ca9caf49 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java @@ -33,7 +33,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -71,7 +71,7 @@ class AuthenticationResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java index 1c13fe8604e..597dc37862e 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java @@ -72,7 +72,7 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -232,7 +232,7 @@ def score( private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private final RedisContainer redis = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer mysql = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer mysql = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer zookeeper = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer clickhouse = ClickHouseContainerUtils.newClickHouseContainer(zookeeper); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java index c4adeb255ad..b12fd0dea91 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java @@ -36,7 +36,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -72,7 +72,7 @@ class ChatCompletionsResourceTest { private static final String USER = RandomStringUtils.randomAlphanumeric(20); private final RedisContainer REDIS_CONTAINER = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MY_SQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MY_SQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetExperimentE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetExperimentE2ETest.java index 0f73fbb9514..5bef0dc4c86 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetExperimentE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetExperimentE2ETest.java @@ -32,8 +32,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; @@ -57,7 +57,7 @@ class DatasetExperimentE2ETest { private static final String EXPERIMENT_RESOURCE_URI = "%s/v1/private/experiments"; private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); private final WireMockUtils.WireMockRuntime wireMock; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java index a882eba8f40..0ceb885e93b 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java @@ -115,8 +115,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; @@ -215,7 +215,7 @@ class DatasetsResourceTest { private static final TimeBasedEpochGenerator GENERATOR = Generators.timeBasedEpochGenerator(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java index 3005951a832..0ff81337a4d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java @@ -109,7 +109,7 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -203,7 +203,7 @@ class ExperimentsResourceTest { private static final TimeBasedEpochGenerator GENERATOR = Generators.timeBasedEpochGenerator(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java index ee3320ad08a..db3159c818c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java @@ -39,8 +39,8 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; @@ -91,7 +91,7 @@ class FeedbackDefinitionResourceTest { private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java index d203e07759d..672a2ed1151 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java @@ -30,7 +30,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -56,7 +56,7 @@ public class GuardrailsResourceTest { private static final String TEST_WORKSPACE = randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java index 0c05f922558..f887806e283 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java @@ -38,7 +38,7 @@ import org.junit.jupiter.params.provider.NullSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -73,7 +73,7 @@ class LlmProviderApiKeyResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); @RegisterApp private final TestDropwizardAppExtension APP; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java index 63ee5d88579..f025c208ace 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java @@ -51,7 +51,7 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -84,7 +84,7 @@ class ManualEvaluationResourceTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); private final GenericContainer ZOOKEEPER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java index 7c52453429d..8b42094c3c5 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java @@ -63,7 +63,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -98,7 +98,7 @@ public class MultiValueFeedbackScoresE2ETest { private static final String EMPTY_REASON_PLACEHOLDER = ""; private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OpenTelemetryResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OpenTelemetryResourceTest.java index cbd165bbf8b..01bc6c389ce 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OpenTelemetryResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OpenTelemetryResourceTest.java @@ -52,8 +52,8 @@ import org.junit.runner.RunWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -94,7 +94,7 @@ class OpenTelemetryResourceTest { public static final String TEST_WORKSPACE = UUID.randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MY_SQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MY_SQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OptimizationsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OptimizationsResourceTest.java index 6acb09f7720..389d81d3cdd 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OptimizationsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/OptimizationsResourceTest.java @@ -51,8 +51,8 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; @@ -89,7 +89,7 @@ class OptimizationsResourceTest { private static final String USER = "user-" + RandomStringUtils.secure().nextAlphanumeric(36); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java index 78c64afa8aa..214167f5046 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java @@ -70,7 +70,7 @@ import org.junit.jupiter.params.provider.NullSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -145,7 +145,7 @@ class ProjectMetricsResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java index 576a26bbc71..27167d51067 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java @@ -74,7 +74,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -149,7 +149,7 @@ class ProjectsResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/PromptResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/PromptResourceTest.java index a3ca2a092e6..b6b14876c80 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/PromptResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/PromptResourceTest.java @@ -58,8 +58,8 @@ import org.junit.jupiter.params.provider.NullSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; @@ -121,7 +121,7 @@ class PromptResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); @RegisterApp private final TestDropwizardAppExtension APP; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java index d57ac00ea05..c710c9d8625 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java @@ -102,7 +102,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -223,7 +223,7 @@ class SpansResourceTest { } private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MY_SQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MY_SQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java index 4a9607840a0..b3019d9ae4c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java @@ -117,8 +117,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; @@ -237,7 +237,7 @@ class TracesResourceTest { } private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java index 927e5e76c23..f99a4f9ef6f 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java @@ -50,7 +50,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -90,7 +90,7 @@ class WorkspacesResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/session/RedirectResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/session/RedirectResourceTest.java index 1ad2537d05a..e455ded0c90 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/session/RedirectResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/session/RedirectResourceTest.java @@ -35,8 +35,8 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; @@ -77,7 +77,7 @@ class RedirectResourceTest { private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE_CONTAINER = ClickHouseContainerUtils .newClickHouseContainer(ZOOKEEPER_CONTAINER); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final WireMockUtils.WireMockRuntime wireMock; @RegisterApp diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java index 7704ed09de3..99e4c9bdf72 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java @@ -16,7 +16,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -31,7 +31,7 @@ class HealthCheckIntegrationTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleCache2E2Test.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleCache2E2Test.java index e2a09a89469..0f7117b443c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleCache2E2Test.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleCache2E2Test.java @@ -23,8 +23,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -50,7 +50,7 @@ class AuthModuleCache2E2Test { private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); private final WireMockUtils.WireMockRuntime wireMock; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleNoAuthIntegrationTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleNoAuthIntegrationTest.java index 91726aa5dc8..158699f8c67 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleNoAuthIntegrationTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/auth/AuthModuleNoAuthIntegrationTest.java @@ -21,8 +21,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -41,7 +41,7 @@ class AuthModuleNoAuthIntegrationTest { public static final String FEEDBACK_DEFINITION_TEMPLATE = "%s/v1/private/feedback-definitions"; private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java index 651428c519f..5a35d77a061 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java @@ -27,7 +27,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.containers.Network; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; @@ -58,7 +58,7 @@ class BiEventListenerTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); private final Network network = Network.newNetwork(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(false); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(false); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(false, network); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(false, network, diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java index 78515ac0485..843a90af69a 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java @@ -38,7 +38,7 @@ import org.quartz.TriggerBuilder; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.containers.Network; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Mono; @@ -211,7 +211,7 @@ INSERT INTO traces ( class CredentialsEnabledScenario { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(false); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(false); private final Network NETWORK = Network.newNetwork(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(false, NETWORK); @@ -355,7 +355,7 @@ private void setUpData(String apiKey, String workspaceName, String workspaceId) class NoCredentialsEnabledScenario { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(false); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(false); private final Network NETWORK = Network.newNetwork(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(false, NETWORK); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/OpikGuiceyLifecycleEventListenerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/OpikGuiceyLifecycleEventListenerTest.java index adfe2eaed6c..d72fb99e91f 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/OpikGuiceyLifecycleEventListenerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/OpikGuiceyLifecycleEventListenerTest.java @@ -22,9 +22,9 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.containers.Network; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.module.lifecycle.GuiceyLifecycle; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -55,7 +55,7 @@ class OpikGuiceyLifecycleEventListenerTest { the reason this is needed is to make sure that only the first test will notify the event. */ - private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(false); + private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(false); private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); private final Network network = Network.newNetwork(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(false, diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java index f68d3f4918b..0447c2a54ad 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java @@ -18,7 +18,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -34,7 +34,7 @@ class CacheManagerTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/health/IsAliveE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/health/IsAliveE2ETest.java index fdfa83913b5..c7c91fff3a4 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/health/IsAliveE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/health/IsAliveE2ETest.java @@ -17,8 +17,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -30,7 +30,7 @@ class IsAliveE2ETest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java index 3db106b8175..6845f2d3d77 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java @@ -17,7 +17,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -31,7 +31,7 @@ class CorsDisabledE2ETest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java index d4898bad966..f27d3d82522 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java @@ -19,7 +19,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; @@ -33,7 +33,7 @@ class CorsEnabledE2ETest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java index 245aa8b0f5b..362daf1a3b5 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java @@ -55,7 +55,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -101,7 +101,7 @@ class RateLimitE2ETest { public static final String TOO_MANY_REQUESTS_MESSAGEE = "Too Many Requests: %s"; private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); private final WireMockUtils.WireMockRuntime wireMock; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/redis/RedissonLockServiceIntegrationTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/redis/RedissonLockServiceIntegrationTest.java index a2cb9bdd959..24ca6712591 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/redis/RedissonLockServiceIntegrationTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/redis/RedissonLockServiceIntegrationTest.java @@ -14,8 +14,8 @@ import org.redisson.api.RedissonReactiveClient; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -40,7 +40,7 @@ class RedissonLockServiceIntegrationTest { private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); - private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER); From 3d9126afafc3312d0458b76de392084d3bad6a5e Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Mon, 3 Nov 2025 10:09:55 +0100 Subject: [PATCH 2/9] Fix imports order --- .../opik/api/resources/v1/events/DatasetEventListenerTest.java | 2 +- .../opik/api/resources/v1/events/OnlineScoringEngineTest.java | 2 +- .../opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java | 2 +- .../com/comet/opik/api/resources/v1/priv/AlertResourceTest.java | 2 +- .../api/resources/v1/priv/AnnotationQueuesResourceTest.java | 2 +- .../opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java | 2 +- .../opik/api/resources/v1/priv/AttachmentResourceTest.java | 2 +- .../opik/api/resources/v1/priv/AuthenticationResourceTest.java | 2 +- .../resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java | 2 +- .../opik/api/resources/v1/priv/ChatCompletionsResourceTest.java | 2 +- .../opik/api/resources/v1/priv/ExperimentsResourceTest.java | 2 +- .../opik/api/resources/v1/priv/GuardrailsResourceTest.java | 2 +- .../api/resources/v1/priv/LlmProviderApiKeyResourceTest.java | 2 +- .../api/resources/v1/priv/ManualEvaluationResourceTest.java | 2 +- .../api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java | 2 +- .../opik/api/resources/v1/priv/ProjectMetricsResourceTest.java | 2 +- .../comet/opik/api/resources/v1/priv/ProjectsResourceTest.java | 2 +- .../com/comet/opik/api/resources/v1/priv/SpansResourceTest.java | 2 +- .../opik/api/resources/v1/priv/WorkspacesResourceTest.java | 2 +- .../comet/opik/infrastructure/HealthCheckIntegrationTest.java | 2 +- .../com/comet/opik/infrastructure/bi/BiEventListenerTest.java | 2 +- .../comet/opik/infrastructure/bi/DailyUsageReportJobTest.java | 2 +- .../com/comet/opik/infrastructure/cache/CacheManagerTest.java | 2 +- .../opik/infrastructure/http/cors/CorsDisabledE2ETest.java | 2 +- .../comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java | 2 +- .../comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java | 2 +- 26 files changed, 26 insertions(+), 26 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java index c53a81d8d61..d62fafbd815 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/DatasetEventListenerTest.java @@ -33,8 +33,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java index 9ce1f9d43ef..164e90e8003 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/events/OnlineScoringEngineTest.java @@ -70,8 +70,8 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java index d84f9ccf47d..cfeaba4ea30 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/jobs/TraceThreadsClosingJobTest.java @@ -35,8 +35,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java index b2a6c75d06f..cda65242c0c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AlertResourceTest.java @@ -73,8 +73,8 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java index 26029c1f022..65e1aa10421 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AnnotationQueuesResourceTest.java @@ -47,8 +47,8 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java index 02c903976fd..280e18b6fb8 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceMinIOTest.java @@ -44,8 +44,8 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java index 8f78ca76da0..d8c2b31fa24 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AttachmentResourceTest.java @@ -36,8 +36,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java index 107ca9caf49..167476eab16 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AuthenticationResourceTest.java @@ -33,8 +33,8 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java index 597dc37862e..2dadab742e3 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java @@ -72,8 +72,8 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java index b12fd0dea91..8c21a1fdd57 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ChatCompletionsResourceTest.java @@ -36,8 +36,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java index 0ff81337a4d..3f5d999244e 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java @@ -109,8 +109,8 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; import ru.vyarus.dropwizard.guice.test.ClientSupport; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java index 672a2ed1151..a655bc18bd5 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GuardrailsResourceTest.java @@ -30,8 +30,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java index f887806e283..bbfb9bd91c6 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java @@ -38,8 +38,8 @@ import org.junit.jupiter.params.provider.NullSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java index f025c208ace..ec8e10cc53c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ManualEvaluationResourceTest.java @@ -51,8 +51,8 @@ import org.mockito.Mockito; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java index 8b42094c3c5..7d46b79dbfe 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/MultiValueFeedbackScoresE2ETest.java @@ -63,8 +63,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java index 214167f5046..c2bfb908f13 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectMetricsResourceTest.java @@ -70,8 +70,8 @@ import org.junit.jupiter.params.provider.NullSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java index 27167d51067..1ae2fa25d7d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java @@ -74,8 +74,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java index c710c9d8625..ade40dec5a6 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java @@ -102,8 +102,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java index f99a4f9ef6f..63ba95f9b56 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/WorkspacesResourceTest.java @@ -50,8 +50,8 @@ import org.junit.jupiter.params.provider.ValueSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java index 99e4c9bdf72..4b85c0f847d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/HealthCheckIntegrationTest.java @@ -16,8 +16,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java index 5a35d77a061..6785b822a7e 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/BiEventListenerTest.java @@ -27,9 +27,9 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.containers.Network; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; import uk.co.jemos.podam.api.PodamFactory; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java index 843a90af69a..db352c7c590 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/bi/DailyUsageReportJobTest.java @@ -38,9 +38,9 @@ import org.quartz.TriggerBuilder; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.containers.Network; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java index 0447c2a54ad..0326fba8944 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/cache/CacheManagerTest.java @@ -18,8 +18,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java index 6845f2d3d77..4256b5e02d6 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsDisabledE2ETest.java @@ -17,8 +17,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java index f27d3d82522..43d138b12f0 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/http/cors/CorsEnabledE2ETest.java @@ -19,8 +19,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import ru.vyarus.dropwizard.guice.test.ClientSupport; import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java index 362daf1a3b5..a5de804c312 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/ratelimit/RateLimitE2ETest.java @@ -55,8 +55,8 @@ import org.junit.jupiter.params.provider.MethodSource; import org.testcontainers.clickhouse.ClickHouseContainer; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.mysql.MySQLContainer; import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import ru.vyarus.dropwizard.guice.test.ClientSupport; From 29811bfd4c02e0cbfd22c8fea39bdabaa34aa4f1 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Mon, 3 Nov 2025 17:11:51 +0100 Subject: [PATCH 3/9] [OPIK-2856] [BE] Implement UUIDv7 time-based filtering for traces - Add InstantToUUIDMapper to convert Instant timestamps to UUIDv7 bounds - Add InstantParamConverter to parse ISO-8601 and epoch millisecond time parameters - Update TracesResource to accept from_time and to_time parameters on /traces and /traces/stats endpoints - Update TraceDAO to apply UUID-based time filtering using BETWEEN clause on id column - Update TraceSearchCriteria to include uuidFromTime and uuidToTime fields - Add comprehensive integration tests for time filtering with boundary conditions - All tests passing: 10/10 time filtering tests + validation tests --- .../java/com/comet/opik/OpikApplication.java | 2 + .../comet/opik/api/InstantToUUIDMapper.java | 51 + .../api/resources/v1/priv/TracesResource.java | 32 +- .../java/com/comet/opik/domain/TraceDAO.java | 16 + .../opik/domain/TraceSearchCriteria.java | 4 +- .../web/InstantParamConverter.java | 61 + .../com/comet/opik/utils/ValidationUtils.java | 16 + .../opik/api/InstantToUUIDMapperTest.java | 224 + .../utils/resources/TraceResourceClient.java | 36 + .../utils/traces/TraceAssertions.java | 32 + .../priv/GetTracesByProjectResourceTest.java | 4996 ++++++++++++++ .../resources/v1/priv/TracesResourceTest.java | 5846 +++-------------- 12 files changed, 6211 insertions(+), 5105 deletions(-) create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java diff --git a/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java b/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java index f32d4f8dae5..b8d7e419a2c 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java @@ -28,6 +28,7 @@ import com.comet.opik.infrastructure.ratelimit.RateLimitModule; import com.comet.opik.infrastructure.redis.RedisModule; import com.comet.opik.infrastructure.usagelimit.UsageLimitModule; +import com.comet.opik.infrastructure.web.InstantParamConverter; import com.comet.opik.utils.JsonBigDecimalDeserializer; import com.comet.opik.utils.JsonUtils; import com.comet.opik.utils.OpenAiMessageJsonDeserializer; @@ -136,5 +137,6 @@ public void run(OpikConfiguration configuration, Environment environment) { jersey.property(ServerProperties.RESPONSE_SET_STATUS_OVER_SEND_ERROR, true); jersey.register(JsonProcessingExceptionMapper.class); + jersey.register(InstantParamConverter.class); } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java new file mode 100644 index 00000000000..279df93845b --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java @@ -0,0 +1,51 @@ +package com.comet.opik.api; + +import com.comet.opik.domain.OpenTelemetryMapper; +import lombok.experimental.UtilityClass; +import lombok.extern.slf4j.Slf4j; + +import java.time.Instant; +import java.util.UUID; + +/** + * Mapper for converting Instant time boundaries to UUIDv7 bounds for efficient BETWEEN queries. + * UUIDv7 encodes the timestamp in the first 48 bits, allowing for lexicographic sorting by time. + */ +@UtilityClass +@Slf4j +public class InstantToUUIDMapper { + + /** + * Generates a UUIDv7 lower bound from a timestamp for BETWEEN queries. + * Uses 0x00 for random bits to get the lexicographically smallest UUID with this timestamp. + * + * @param timestamp the instant in time + * @return the lower bound UUIDv7 + */ + public static UUID toLowerBound(Instant timestamp) { + if (timestamp == null) { + return null; + } + + return OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], timestamp.toEpochMilli()); + } + + /** + * Generates a UUIDv7 upper bound from a timestamp for BETWEEN queries. + * Uses zero bytes for the NEXT millisecond to get the first UUID after the end time. + * This ensures BETWEEN includes all UUIDs created within the end timestamp. + * + * @param timestamp the instant in time + * @return the upper bound UUIDv7 (UUID for next millisecond with zero random bits) + */ + public static UUID toUpperBound(Instant timestamp) { + if (timestamp == null) { + return null; + } + + // Add 1ms and use zero random bytes to get the first UUID AFTER the end time + // BETWEEN will include all UUIDs from toLowerBound to (but not including) this value + long nextMillis = timestamp.toEpochMilli() + 1; + return OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], nextMillis); + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java index 46a73a224a8..8f403b6ee52 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java @@ -11,6 +11,7 @@ import com.comet.opik.api.FeedbackScore; import com.comet.opik.api.FeedbackScoreBatchContainer; import com.comet.opik.api.FeedbackScoreNames; +import com.comet.opik.api.InstantToUUIDMapper; import com.comet.opik.api.ProjectStats; import com.comet.opik.api.Trace; import com.comet.opik.api.Trace.TracePage; @@ -80,6 +81,7 @@ import org.glassfish.jersey.server.ChunkedOutput; import reactor.core.publisher.Flux; +import java.time.Instant; import java.util.List; import java.util.Optional; import java.util.Set; @@ -90,6 +92,7 @@ import static com.comet.opik.api.TraceThread.TraceThreadPage; import static com.comet.opik.utils.AsyncUtils.setRequestContext; import static com.comet.opik.utils.ValidationUtils.validateProjectNameAndProjectId; +import static com.comet.opik.utils.ValidationUtils.validateTimeRangeParameters; @Path("/v1/private/traces") @Produces(MediaType.APPLICATION_JSON) @@ -125,9 +128,12 @@ public Response getTracesByProject( @QueryParam("truncate") @DefaultValue("false") @Schema(description = "Truncate input, output and metadata to slim payloads") boolean truncate, @QueryParam("strip_attachments") @DefaultValue("false") @Schema(description = "If true, returns attachment references like [file.png]; if false, downloads and reinjects stripped attachments") boolean stripAttachments, @QueryParam("sorting") String sorting, - @QueryParam("exclude") String exclude) { + @QueryParam("exclude") String exclude, + @QueryParam("from_time") Instant startTime, + @QueryParam("to_time") Instant endTime) { validateProjectNameAndProjectId(projectName, projectId); + validateTimeRangeParameters(startTime, endTime); var traceFilters = filtersFactory.newFilters(filters, TraceFilter.LIST_TYPE_REFERENCE); var sortingFields = traceSortingFactory.newSorting(sorting); @@ -149,8 +155,10 @@ public Response getTracesByProject( .filters(traceFilters) .truncate(truncate) .stripAttachments(stripAttachments) - .sortingFields(sortingFields) + .uuidFromTime(InstantToUUIDMapper.toLowerBound(startTime)) + .uuidToTime(InstantToUUIDMapper.toUpperBound(endTime)) .exclude(ParamsValidator.get(exclude, Trace.TraceField.class, "exclude")) + .sortingFields(sortingFields) .build(); log.info("Get traces by '{}' on workspaceId '{}'", searchCriteria, workspaceId); @@ -353,15 +361,27 @@ public Response deleteTraces( @JsonView({ProjectStats.ProjectStatItem.View.Public.class}) public Response getStats(@QueryParam("project_id") UUID projectId, @QueryParam("project_name") String projectName, - @QueryParam("filters") String filters) { + @QueryParam("filters") String filters, + @QueryParam("from_time") Instant startTime, + @QueryParam("to_time") Instant endTime) { validateProjectNameAndProjectId(projectName, projectId); + validateTimeRangeParameters(startTime, endTime); var traceFilters = filtersFactory.newFilters(filters, TraceFilter.LIST_TYPE_REFERENCE); - var searchCriteria = TraceSearchCriteria.builder() + + var searchCriteriaBuilder = TraceSearchCriteria.builder() .projectName(projectName) .projectId(projectId) - .filters(traceFilters) - .build(); + .filters(traceFilters); + + // Apply UUID-based time filtering if both from_time and to_time are provided + if (startTime != null && endTime != null) { + searchCriteriaBuilder + .uuidFromTime(InstantToUUIDMapper.toLowerBound(startTime)) + .uuidToTime(InstantToUUIDMapper.toUpperBound(endTime)); + } + + var searchCriteria = searchCriteriaBuilder.build(); String workspaceId = requestContext.get().getWorkspaceId(); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java index 5e8e1259da0..194280873bb 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java @@ -781,6 +781,7 @@ AND notEquals(start_time, toDateTime64('1970-01-01 00:00:00.000', 9)), WHERE workspace_id = :workspace_id AND project_id = :project_id + AND id BETWEEN :uuid_from_time AND :uuid_to_time AND id \\< :last_received_id AND AND @@ -995,6 +996,7 @@ AND notEquals(start_time, toDateTime64('1970-01-01 00:00:00.000', 9)), WHERE project_id = :project_id AND workspace_id = :workspace_id + AND id BETWEEN :uuid_from_time AND :uuid_to_time AND AND @@ -1412,6 +1414,7 @@ AND notEquals(start_time, toDateTime64('1970-01-01 00:00:00.000', 9)), WHERE workspace_id = :workspace_id AND project_id IN :project_ids + AND id BETWEEN :uuid_from_time AND :uuid_to_time AND AND @@ -2892,6 +2895,13 @@ private ST newFindTemplate(String query, TraceSearchCriteria traceSearchCriteria }); Optional.ofNullable(traceSearchCriteria.lastReceivedId()) .ifPresent(lastReceivedTraceId -> template.add("last_received_id", lastReceivedTraceId)); + + // Bind UUID BETWEEN bounds for time-based filtering + if (traceSearchCriteria.uuidFromTime() != null && traceSearchCriteria.uuidToTime() != null) { + template.add("uuid_between", true); + template.add("uuid_from_time", traceSearchCriteria.uuidFromTime()); + template.add("uuid_to_time", traceSearchCriteria.uuidToTime()); + } return template; } @@ -2907,6 +2917,12 @@ private void bindSearchCriteria(TraceSearchCriteria traceSearchCriteria, Stateme }); Optional.ofNullable(traceSearchCriteria.lastReceivedId()) .ifPresent(lastReceivedTraceId -> statement.bind("last_received_id", lastReceivedTraceId)); + + // Bind UUID BETWEEN bounds for time-based filtering + if (traceSearchCriteria.uuidFromTime() != null && traceSearchCriteria.uuidToTime() != null) { + statement.bind("uuid_from_time", traceSearchCriteria.uuidFromTime()); + statement.bind("uuid_to_time", traceSearchCriteria.uuidToTime()); + } } @Override diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceSearchCriteria.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceSearchCriteria.java index af1c5538bd5..7fee3e1847d 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceSearchCriteria.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceSearchCriteria.java @@ -18,5 +18,7 @@ public record TraceSearchCriteria( UUID lastReceivedId, boolean truncate, boolean stripAttachments, - Set exclude) { + Set exclude, + UUID uuidFromTime, + UUID uuidToTime) { } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java new file mode 100644 index 00000000000..32846054474 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java @@ -0,0 +1,61 @@ +package com.comet.opik.infrastructure.web; + +import jakarta.ws.rs.BadRequestException; +import jakarta.ws.rs.ext.ParamConverter; +import jakarta.ws.rs.ext.ParamConverterProvider; +import jakarta.ws.rs.ext.Provider; +import lombok.extern.slf4j.Slf4j; + +import java.lang.annotation.Annotation; +import java.lang.reflect.Type; +import java.time.Instant; + +/** + * JAX-RS ParamConverter for automatic Instant conversion from query parameters. + * Supports both ISO-8601 format (e.g., "2024-01-01T00:00:00Z") and milliseconds since epoch. + */ +@Provider +@Slf4j +public class InstantParamConverter implements ParamConverterProvider { + + @Override + public ParamConverter getConverter(Class rawType, Type genericType, Annotation[] annotations) { + if (rawType != Instant.class) { + return null; + } + + return (ParamConverter) new ParamConverter() { + @Override + public Instant fromString(String value) { + if (value == null || value.isEmpty()) { + return null; + } + + try { + // Try parsing as ISO-8601 format first + return Instant.parse(value); + } catch (Exception exception) { + log.debug("Failed to parse '{}' as ISO-8601, attempting to parse as milliseconds since epoch", + value); + try { + // Fall back to parsing as milliseconds since epoch + long epochMillis = Long.parseLong(value); + return Instant.ofEpochMilli(epochMillis); + } catch (NumberFormatException numberFormatException) { + log.error( + "Invalid instant format: '{}'. Expected ISO-8601 (e.g., 2024-01-01T00:00:00Z) or milliseconds since epoch", + value); + throw new BadRequestException( + "Invalid instant format: '%s'. Expected ISO-8601 (e.g., 2024-01-01T00:00:00Z) or milliseconds since epoch" + .formatted(value)); + } + } + } + + @Override + public String toString(Instant value) { + return value != null ? value.toString() : null; + } + }; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java index b0bc787266e..eef70217c77 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java @@ -3,6 +3,7 @@ import jakarta.ws.rs.BadRequestException; import org.apache.commons.lang3.StringUtils; +import java.time.Instant; import java.util.UUID; public class ValidationUtils { @@ -49,4 +50,19 @@ public static void validateProjectNameAndProjectId(String projectName, UUID proj throw new BadRequestException("Either 'project_name' or 'project_id' query params must be provided"); } } + + public static void validateTimeRangeParameters(Instant startTime, Instant endTime) { + boolean startTimePresent = startTime != null; + boolean endTimePresent = endTime != null; + + if (startTimePresent != endTimePresent) { + throw new BadRequestException( + "Both 'from_time' and 'to_time' parameters must be provided together, or both must be omitted"); + } + + if (startTimePresent && endTimePresent && startTime.isAfter(endTime)) { + throw new BadRequestException( + "Parameter 'from_time' must be before 'to_time'"); + } + } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java new file mode 100644 index 00000000000..f3c83f5f6df --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java @@ -0,0 +1,224 @@ +package com.comet.opik.api; + +import com.comet.opik.domain.OpenTelemetryMapper; +import org.junit.jupiter.api.Test; + +import java.time.Instant; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for InstantToUUIDMapper to validate UUIDv7 boundary generation + * and ensure consistent timestamp encoding. + */ +class InstantToUUIDMapperTest { + + @Test + void shouldGenerateLowerBound_withZeroBytesForRandomPortion() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + + // Then + assertThat(lowerBound).isNotNull(); + assertThat(lowerBound.toString()).matches("[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[0-9a-f]{4}-[0-9a-f]{12}"); + } + + @Test + void shouldGenerateUpperBound_withAllFFBytesForRandomPortion() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + + // Then + assertThat(upperBound).isNotNull(); + assertThat(upperBound.toString()).matches("[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[0-9a-f]{4}-[0-9a-f]{12}"); + } + + @Test + void shouldProduceDifferentUUIDsForZeroAndFFBytes() { + // Given - Despite trying to use 0x00 and 0xFF for random bytes, + // they get SHA-256 hashed, so we just verify they produce different UUIDs + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + + // Then + assertThat(lowerBound).isNotEqualTo(upperBound); + } + + @Test + void shouldEncodeSameTimestampInBothBounds() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + + // Then + // Extract the timestamp portion (first 48 bits / first 12 hex chars) + String lowerBoundStr = lowerBound.toString().replace("-", ""); + String upperBoundStr = upperBound.toString().replace("-", ""); + + String lowerTimestampPart = lowerBoundStr.substring(0, 12); + String upperTimestampPart = upperBoundStr.substring(0, 12); + + assertThat(lowerTimestampPart) + .as("Both bounds should encode the same timestamp") + .isEqualTo(upperTimestampPart); + } + + @Test + void shouldHaveDifferentRandomPortions() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + + // Then + // Extract the random portion (after the timestamp and version) + String lowerBoundStr = lowerBound.toString().replace("-", ""); + String upperBoundStr = upperBound.toString().replace("-", ""); + + String lowerRandomPart = lowerBoundStr.substring(13); // Skip timestamp+version + String upperRandomPart = upperBoundStr.substring(13); + + assertThat(lowerRandomPart) + .as("Random portions should differ (lower should be 0000, upper should be FFFF)") + .isNotEqualTo(upperRandomPart); + } + + @Test + void shouldMaintainChronologicalOrder() { + // Given + Instant time1 = Instant.parse("2025-01-15T10:00:00Z"); + Instant time2 = Instant.parse("2025-01-15T11:00:00Z"); + Instant time3 = Instant.parse("2025-01-15T12:00:00Z"); + + // When + UUID lower1 = InstantToUUIDMapper.toLowerBound(time1); + UUID lower2 = InstantToUUIDMapper.toLowerBound(time2); + UUID lower3 = InstantToUUIDMapper.toLowerBound(time3); + + // Then + assertThat(lower1.compareTo(lower2)).isNegative(); + assertThat(lower2.compareTo(lower3)).isNegative(); + assertThat(lower1.compareTo(lower3)).isNegative(); + } + + @Test + void shouldReturnNull_whenTimestampIsNull() { + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(null); + UUID upperBound = InstantToUUIDMapper.toUpperBound(null); + + // Then + assertThat(lowerBound).isNull(); + assertThat(upperBound).isNull(); + } + + @Test + void shouldProduceUUIDv7Format() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + + // Then + // UUIDv7 has version nibble = 7 in the 13th character + String uuidStr = lowerBound.toString(); + assertThat(uuidStr.charAt(14)).isEqualTo('7'); + } + + @Test + void shouldGenerateConsistentUUIDs() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound1 = InstantToUUIDMapper.toLowerBound(timestamp); + UUID lowerBound2 = InstantToUUIDMapper.toLowerBound(timestamp); + + // Then + assertThat(lowerBound1).isEqualTo(lowerBound2); + } + + @Test + void shouldUseZeroBytesForLowerBound() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + UUID referenceUUID = OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], timestamp.toEpochMilli()); + + // Then + assertThat(lowerBound).isEqualTo(referenceUUID); + } + + @Test + void shouldUseFFBytesForUpperBound() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + UUID referenceUUID = OpenTelemetryMapper.convertOtelIdToUUIDv7( + new byte[]{-1, -1, -1, -1, -1, -1, -1, -1}, + timestamp.toEpochMilli()); + + // Then + assertThat(upperBound).isEqualTo(referenceUUID); + } + + @Test + void shouldWorkWithBoundaryInstants() { + // Given - Test with epoch start + Instant epochStart = Instant.EPOCH; + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(epochStart); + UUID upperBound = InstantToUUIDMapper.toUpperBound(epochStart); + + // Then + assertThat(lowerBound).isNotNull(); + assertThat(upperBound).isNotNull(); + // They should be different due to different random bytes (hashed) + assertThat(lowerBound).isNotEqualTo(upperBound); + } + + @Test + void shouldAllUUIDsWithSameTimestampShareTimestampPortion() { + // Given + Instant queryTime = Instant.parse("2025-01-15T10:30:00Z"); + UUID lowerBound = InstantToUUIDMapper.toLowerBound(queryTime); + UUID upperBound = InstantToUUIDMapper.toUpperBound(queryTime); + UUID randomUUID = OpenTelemetryMapper.convertOtelIdToUUIDv7( + new byte[]{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}, + queryTime.toEpochMilli()); + + // Then - All UUIDs with the same timestamp should have the same timestamp portion + String lowerStr = lowerBound.toString().replace("-", ""); + String upperStr = upperBound.toString().replace("-", ""); + String randomStr = randomUUID.toString().replace("-", ""); + + // Extract timestamp (first 12 hex chars = 48 bits) + String lowerTimestamp = lowerStr.substring(0, 12); + String upperTimestamp = upperStr.substring(0, 12); + String randomTimestamp = randomStr.substring(0, 12); + + assertThat(lowerTimestamp).isEqualTo(upperTimestamp); + assertThat(lowerTimestamp).isEqualTo(randomTimestamp); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/TraceResourceClient.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/TraceResourceClient.java index 5004d076a3d..eb0ae8ed467 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/TraceResourceClient.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/TraceResourceClient.java @@ -520,6 +520,13 @@ public Trace.TracePage getTraces(String projectName, UUID projectId, String apiK target = target.queryParam("project_id", projectId); } + // Add remaining queryParams (like from_time, to_time) + WebTarget finalTarget = target; + target = queryParams.entrySet() + .stream() + .filter(e -> !e.getKey().equals("page")) // Skip page as it's already handled + .reduce(finalTarget, (acc, entry) -> acc.queryParam(entry.getKey(), entry.getValue()), (a, b) -> b); + var actualResponse = target .queryParam("filters", toURLEncodedQueryParam(filters)) .request() @@ -660,4 +667,33 @@ public Response callBatchCreateTracesWithCookie(List traces, String sessi .header(WORKSPACE_HEADER, workspaceName) .post(Entity.json(new TraceBatch(traces))); } + + public ProjectStats getStats(String projectName, UUID projectId, String apiKey, String workspaceName, + Map queryParams) { + WebTarget webTarget = client.target(RESOURCE_PATH.formatted(baseURI)) + .path("stats"); + + if (projectName != null) { + webTarget = webTarget.queryParam("project_name", projectName); + } + if (projectId != null) { + webTarget = webTarget.queryParam("project_id", projectId); + } + + // Add remaining queryParams (like from_time, to_time) + WebTarget finalTarget = webTarget; + webTarget = queryParams.entrySet() + .stream() + .reduce(finalTarget, (acc, entry) -> acc.queryParam(entry.getKey(), entry.getValue()), (a, b) -> b); + + try (var actualResponse = webTarget + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + assertThat(actualResponse.getStatus()).isEqualTo(HttpStatus.SC_OK); + return actualResponse.readEntity(ProjectStats.class); + } + } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/traces/TraceAssertions.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/traces/TraceAssertions.java index e466868b7ee..f2769d6357b 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/traces/TraceAssertions.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/traces/TraceAssertions.java @@ -11,7 +11,10 @@ import java.math.BigDecimal; import java.time.Instant; import java.time.temporal.ChronoUnit; +import java.util.EnumMap; import java.util.List; +import java.util.Map; +import java.util.function.Function; import static com.comet.opik.api.resources.utils.CommentAssertionUtils.assertComments; import static org.assertj.core.api.Assertions.assertThat; @@ -30,6 +33,35 @@ public class TraceAssertions { "threadModelId", "feedbackScores.createdAt", "feedbackScores.lastUpdatedAt", "feedbackScores.valueByAuthor"}; + public static final Map> EXCLUDE_FUNCTIONS = new EnumMap<>( + Trace.TraceField.class); + + static { + EXCLUDE_FUNCTIONS.put(Trace.TraceField.NAME, it -> it.toBuilder().name(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.START_TIME, it -> it.toBuilder().startTime(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.END_TIME, it -> it.toBuilder().endTime(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.INPUT, it -> it.toBuilder().input(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.OUTPUT, it -> it.toBuilder().output(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.METADATA, it -> it.toBuilder().metadata(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.TAGS, it -> it.toBuilder().tags(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.USAGE, it -> it.toBuilder().usage(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.ERROR_INFO, it -> it.toBuilder().errorInfo(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.CREATED_AT, it -> it.toBuilder().createdAt(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.CREATED_BY, it -> it.toBuilder().createdBy(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.LAST_UPDATED_BY, it -> it.toBuilder().lastUpdatedBy(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.FEEDBACK_SCORES, it -> it.toBuilder().feedbackScores(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.COMMENTS, it -> it.toBuilder().comments(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.GUARDRAILS_VALIDATIONS, + it -> it.toBuilder().guardrailsValidations(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.SPAN_COUNT, it -> it.toBuilder().spanCount(0).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.LLM_SPAN_COUNT, it -> it.toBuilder().llmSpanCount(0).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.TOTAL_ESTIMATED_COST, + it -> it.toBuilder().totalEstimatedCost(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.THREAD_ID, it -> it.toBuilder().threadId(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.DURATION, it -> it.toBuilder().duration(null).build()); + EXCLUDE_FUNCTIONS.put(Trace.TraceField.VISIBILITY_MODE, it -> it.toBuilder().visibilityMode(null).build()); + } + public static void assertErrorResponse(Response actualResponse, String message, int expectedStatus) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedStatus); assertThat(actualResponse.hasEntity()).isTrue(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java new file mode 100644 index 00000000000..d61283e991b --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java @@ -0,0 +1,4996 @@ +package com.comet.opik.api.resources.v1.priv; + +import com.comet.opik.api.AnnotationQueue; +import com.comet.opik.api.Comment; +import com.comet.opik.api.ErrorInfo; +import com.comet.opik.api.FeedbackScore; +import com.comet.opik.api.FeedbackScoreItem.FeedbackScoreBatchItem.FeedbackScoreBatchItemBuilder; +import com.comet.opik.api.Guardrail; +import com.comet.opik.api.Project; +import com.comet.opik.api.Span; +import com.comet.opik.api.Trace; +import com.comet.opik.api.TraceSearchStreamRequest; +import com.comet.opik.api.VisibilityMode; +import com.comet.opik.api.filter.Field; +import com.comet.opik.api.filter.FieldType; +import com.comet.opik.api.filter.Filter; +import com.comet.opik.api.filter.Operator; +import com.comet.opik.api.filter.TraceField; +import com.comet.opik.api.filter.TraceFilter; +import com.comet.opik.api.resources.utils.AuthTestUtils; +import com.comet.opik.api.resources.utils.ClickHouseContainerUtils; +import com.comet.opik.api.resources.utils.ClientSupportUtils; +import com.comet.opik.api.resources.utils.DurationUtils; +import com.comet.opik.api.resources.utils.MigrationUtils; +import com.comet.opik.api.resources.utils.MinIOContainerUtils; +import com.comet.opik.api.resources.utils.MySQLContainerUtils; +import com.comet.opik.api.resources.utils.RedisContainerUtils; +import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; +import com.comet.opik.api.resources.utils.WireMockUtils; +import com.comet.opik.api.resources.utils.resources.AnnotationQueuesResourceClient; +import com.comet.opik.api.resources.utils.resources.GuardrailsGenerator; +import com.comet.opik.api.resources.utils.resources.GuardrailsResourceClient; +import com.comet.opik.api.resources.utils.resources.ProjectResourceClient; +import com.comet.opik.api.resources.utils.resources.SpanResourceClient; +import com.comet.opik.api.resources.utils.resources.TraceResourceClient; +import com.comet.opik.api.resources.utils.traces.TraceAssertions; +import com.comet.opik.api.resources.utils.traces.TracePageTestAssertion; +import com.comet.opik.api.resources.utils.traces.TraceStatsAssertion; +import com.comet.opik.api.resources.utils.traces.TraceStreamTestAssertion; +import com.comet.opik.api.resources.utils.traces.TraceTestAssertion; +import com.comet.opik.api.sorting.Direction; +import com.comet.opik.api.sorting.SortableFields; +import com.comet.opik.api.sorting.SortingField; +import com.comet.opik.domain.GuardrailResult; +import com.comet.opik.domain.GuardrailsMapper; +import com.comet.opik.domain.OpenTelemetryMapper; +import com.comet.opik.domain.SpanType; +import com.comet.opik.domain.cost.CostService; +import com.comet.opik.domain.filter.FilterQueryBuilder; +import com.comet.opik.extensions.DropwizardAppExtensionProvider; +import com.comet.opik.extensions.RegisterApp; +import com.comet.opik.podam.PodamFactoryUtils; +import com.comet.opik.utils.JsonUtils; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.uuid.Generators; +import com.fasterxml.uuid.impl.TimeBasedEpochGenerator; +import com.google.common.collect.Lists; +import com.redis.testcontainers.RedisContainer; +import jakarta.ws.rs.client.Entity; +import jakarta.ws.rs.client.WebTarget; +import jakarta.ws.rs.core.HttpHeaders; +import jakarta.ws.rs.core.Response; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.lang3.RandomStringUtils; +import org.apache.commons.lang3.RandomUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.http.HttpStatus; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.testcontainers.clickhouse.ClickHouseContainer; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; +import ru.vyarus.dropwizard.guice.test.ClientSupport; +import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; +import uk.co.jemos.podam.api.PodamFactory; +import uk.co.jemos.podam.api.PodamUtils; + +import java.math.BigDecimal; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Comparator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static com.comet.opik.api.FeedbackScoreItem.FeedbackScoreBatchItem; +import static com.comet.opik.api.filter.TraceField.CUSTOM; +import static com.comet.opik.api.resources.utils.ClickHouseContainerUtils.DATABASE_NAME; +import static com.comet.opik.api.resources.utils.TestUtils.toURLEncodedQueryParam; +import static com.comet.opik.api.resources.utils.traces.TraceAssertions.IGNORED_FIELDS_TRACES; +import static com.comet.opik.infrastructure.auth.RequestContext.WORKSPACE_HEADER; +import static java.util.UUID.randomUUID; +import static java.util.stream.Collectors.toCollection; +import static java.util.stream.Collectors.toMap; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +@DisplayName("Get Traces Resource Test") +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@ExtendWith(DropwizardAppExtensionProvider.class) +class GetTracesByProjectResourceTest { + + public static final String URL_TEMPLATE = "%s/v1/private/traces"; + + private static final String API_KEY = UUID.randomUUID().toString(); + private static final String USER = UUID.randomUUID().toString(); + private static final String WORKSPACE_ID = UUID.randomUUID().toString(); + private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); + + private final RedisContainer redis = RedisContainerUtils.newRedisContainer(); + private final MySQLContainer mysqlContainer = MySQLContainerUtils.newMySQLContainer(); + private final GenericContainer zookeeperContainer = ClickHouseContainerUtils.newZookeeperContainer(); + private final ClickHouseContainer clickHouseContainer = ClickHouseContainerUtils + .newClickHouseContainer(zookeeperContainer); + private final GenericContainer minio = MinIOContainerUtils.newMinIOContainer(); + + private final WireMockUtils.WireMockRuntime wireMock; + + @RegisterApp + private final TestDropwizardAppExtension app; + + { + Startables.deepStart(redis, mysqlContainer, clickHouseContainer, zookeeperContainer, minio).join(); + String minioUrl = "http://%s:%d".formatted(minio.getHost(), minio.getMappedPort(9000)); + + wireMock = WireMockUtils.startWireMock(); + + var databaseAnalyticsFactory = ClickHouseContainerUtils.newDatabaseAnalyticsFactory( + clickHouseContainer, DATABASE_NAME); + + MigrationUtils.runMysqlDbMigration(mysqlContainer); + MigrationUtils.runClickhouseDbMigration(clickHouseContainer); + MinIOContainerUtils.setupBucketAndCredentials(minioUrl); + + app = TestDropwizardAppExtensionUtils.newTestDropwizardAppExtension( + TestDropwizardAppExtensionUtils.AppContextConfig.builder() + .jdbcUrl(mysqlContainer.getJdbcUrl()) + .databaseAnalyticsFactory(databaseAnalyticsFactory) + .redisUrl(redis.getRedisURI()) + .runtimeInfo(wireMock.runtimeInfo()) + .isMinIO(true) + .minioUrl(minioUrl) + .build()); + } + + private final PodamFactory factory = PodamFactoryUtils.newPodamFactory(); + private final TimeBasedEpochGenerator generator = Generators.timeBasedEpochGenerator(); + private final FilterQueryBuilder filterQueryBuilder = new FilterQueryBuilder(); + + private String baseURI; + private ClientSupport client; + private ProjectResourceClient projectResourceClient; + private TraceResourceClient traceResourceClient; + private SpanResourceClient spanResourceClient; + private GuardrailsResourceClient guardrailsResourceClient; + private GuardrailsGenerator guardrailsGenerator; + private AnnotationQueuesResourceClient annotationQueuesResourceClient; + + @BeforeAll + void setUpAll(ClientSupport client) { + + this.baseURI = TestUtils.getBaseUrl(client); + this.client = client; + + ClientSupportUtils.config(client); + + mockTargetWorkspace(API_KEY, TEST_WORKSPACE, WORKSPACE_ID); + + this.projectResourceClient = new ProjectResourceClient(this.client, baseURI, factory); + this.traceResourceClient = new TraceResourceClient(this.client, baseURI); + this.spanResourceClient = new SpanResourceClient(this.client, baseURI); + this.guardrailsResourceClient = new GuardrailsResourceClient(client, baseURI); + this.annotationQueuesResourceClient = new AnnotationQueuesResourceClient(client, baseURI); + this.guardrailsGenerator = new GuardrailsGenerator(); + } + + private void mockTargetWorkspace(String apiKey, String workspaceName, String workspaceId) { + AuthTestUtils.mockTargetWorkspace(wireMock.server(), apiKey, workspaceName, workspaceId, USER); + } + + @AfterAll + void tearDownAll() { + wireMock.server().stop(); + } + + private UUID getProjectId(String projectName, String workspaceName, String apiKey) { + return projectResourceClient.getByName(projectName, apiKey, workspaceName).id(); + } + + @Nested + @DisplayName("Filters Test:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class FilterTest { + + private final TraceStatsAssertion traceStatsAssertion = new TraceStatsAssertion(traceResourceClient); + private final TraceTestAssertion traceTestAssertion = new TraceTestAssertion(traceResourceClient, USER); + private final TraceStreamTestAssertion traceStreamTestAssertion = new TraceStreamTestAssertion( + traceResourceClient, USER); + + private Stream getFilterTestArguments() { + return Stream.of( + Arguments.of( + "/traces/stats", + traceStatsAssertion), + Arguments.of( + "/traces", + traceTestAssertion), + Arguments.of( + "/traces/search", + traceStreamTestAssertion)); + } + + private Stream equalAndNotEqualFilters() { + return Stream.of( + Arguments.of( + "/traces/stats", + Operator.EQUAL, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceStatsAssertion), + Arguments.of( + "/traces", + Operator.EQUAL, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceTestAssertion), + Arguments.of( + "/traces/search", + Operator.EQUAL, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceStreamTestAssertion), + Arguments.of( + "/traces/stats", + Operator.NOT_EQUAL, + (Function, List>) traces -> traces.subList(1, traces.size()), + (Function, List>) traces -> List.of(traces.getFirst()), + traceStatsAssertion), + Arguments.of( + "/traces", + Operator.NOT_EQUAL, + (Function, List>) traces -> traces.subList(1, traces.size()), + (Function, List>) traces -> List.of(traces.getFirst()), + traceTestAssertion), + Arguments.of( + "/traces/search", + Operator.NOT_EQUAL, + (Function, List>) traces -> traces.subList(1, traces.size()), + (Function, List>) traces -> List.of(traces.getFirst()), + traceStreamTestAssertion)); + } + + private Stream getUsageKeyArgs() { + return Stream.of( + Arguments.of( + "/traces/stats", + traceStatsAssertion, + "completion_tokens", + TraceField.USAGE_COMPLETION_TOKENS), + Arguments.of( + "/traces/stats", + traceStatsAssertion, + "prompt_tokens", + TraceField.USAGE_PROMPT_TOKENS), + Arguments.of( + "/traces/stats", + traceStatsAssertion, + "total_tokens", + TraceField.USAGE_TOTAL_TOKENS), + Arguments.of( + "/traces", + traceTestAssertion, + "completion_tokens", + TraceField.USAGE_COMPLETION_TOKENS), + Arguments.of( + "/traces", + traceTestAssertion, + "prompt_tokens", + TraceField.USAGE_PROMPT_TOKENS), + Arguments.of( + "/traces", + traceTestAssertion, + "total_tokens", + TraceField.USAGE_TOTAL_TOKENS), + Arguments.of( + "/traces/search", + traceStreamTestAssertion, + "completion_tokens", + TraceField.USAGE_COMPLETION_TOKENS), + Arguments.of( + "/traces/search", + traceStreamTestAssertion, + "prompt_tokens", + TraceField.USAGE_PROMPT_TOKENS), + Arguments.of( + "/traces/search", + traceStreamTestAssertion, + "total_tokens", + TraceField.USAGE_TOTAL_TOKENS)); + } + + private Stream getFeedbackScoresArgs() { + return Stream.of( + Arguments.of( + "/traces/stats", + Operator.EQUAL, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceStatsAssertion), + Arguments.of( + "/traces", + Operator.EQUAL, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceTestAssertion), + Arguments.of( + "/traces/search", + Operator.EQUAL, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceStreamTestAssertion), + Arguments.of( + "/traces/stats", + Operator.NOT_EQUAL, + (Function, List>) traces -> traces.subList(2, traces.size()), + (Function, List>) traces -> traces.subList(0, 2), + traceStatsAssertion), + Arguments.of( + "/traces", + Operator.NOT_EQUAL, + (Function, List>) traces -> traces.subList(2, traces.size()), + (Function, List>) traces -> traces.subList(0, 2), + traceTestAssertion), + Arguments.of( + "/traces/search", + Operator.NOT_EQUAL, + (Function, List>) traces -> traces.subList(2, traces.size()), + (Function, List>) traces -> traces.subList(0, 2), + traceStreamTestAssertion)); + } + + private Stream getDurationArgs() { + Stream arguments = Stream.of( + arguments(Operator.EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 1.0), + arguments(Operator.GREATER_THAN, Duration.ofMillis(8L).toNanos() / 1000, 7.0), + arguments(Operator.GREATER_THAN_EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 1.0), + arguments(Operator.GREATER_THAN_EQUAL, Duration.ofMillis(1L).plusNanos(1000).toNanos() / 1000, 1.0), + arguments(Operator.LESS_THAN, Duration.ofMillis(1L).plusNanos(1).toNanos() / 1000, 2.0), + arguments(Operator.LESS_THAN_EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 1.0), + arguments(Operator.LESS_THAN_EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 2.0)); + + return arguments.flatMap(arg -> Stream.of( + arguments("/traces/stats", traceStatsAssertion, arg.get()[0], + arg.get()[1], arg.get()[2]), + arguments("/traces", traceTestAssertion, arg.get()[0], + arg.get()[1], arg.get()[2]), + arguments("/traces/search", traceStreamTestAssertion, + arg.get()[0], + arg.get()[1], arg.get()[2]))); + } + + private Stream getFilterInvalidOperatorForFieldTypeArgs() { + return filterQueryBuilder.getUnSupportedOperators(TraceField.values()) + .entrySet() + .stream() + .flatMap(filter -> filter.getValue() + .stream() + .flatMap(operator -> Stream.of( + Arguments.of("/stats", TraceFilter.builder() + .field(filter.getKey()) + .operator(operator) + .key(getKey(filter.getKey())) + .value(getValidValue(filter.getKey())) + .build()), + Arguments.of("/search", TraceFilter.builder() + .field(filter.getKey()) + .operator(operator) + .key(getKey(filter.getKey())) + .value(getValidValue(filter.getKey())) + .build()), + Arguments.of("", TraceFilter.builder() + .field(filter.getKey()) + .operator(operator) + .key(getKey(filter.getKey())) + .value(getValidValue(filter.getKey())) + .build())))); + } + + private Stream getFilterInvalidValueOrKeyForFieldTypeArgs() { + + Stream filters = filterQueryBuilder.getSupportedOperators(TraceField.values()) + .entrySet() + .stream() + .flatMap(filter -> filter.getValue() + .stream() + .flatMap(operator -> switch (filter.getKey().getType()) { + case STRING -> Stream.empty(); + case DICTIONARY, FEEDBACK_SCORES_NUMBER -> Stream.of( + TraceFilter.builder() + .field(filter.getKey()) + .operator(operator) + .key(null) + .value(getValidValue(filter.getKey())) + .build(), + TraceFilter.builder() + .field(filter.getKey()) + .operator(operator) + // if no value is expected, create an invalid filter by an empty key + .key(Operator.NO_VALUE_OPERATORS.contains(operator) + ? "" + : getKey(filter.getKey())) + .value(getInvalidValue(filter.getKey())) + .build()); + case ERROR_CONTAINER -> Stream.of(); + default -> Stream.of(TraceFilter.builder() + .field(filter.getKey()) + .operator(operator) + .value(getInvalidValue(filter.getKey())) + .build()); + })); + + return filters.flatMap(filter -> Stream.of( + arguments("/stats", filter), + arguments("", filter), + arguments("/search", filter))); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + @DisplayName("when project name and project id are null, then return bad request") + void whenProjectNameAndIdAreNull__thenReturnBadRequest(String endpoint, TracePageTestAssertion testAssertion) { + + Project project = factory.manufacturePojo(Project.class); + var projectId = projectResourceClient.createProject(project, API_KEY, TEST_WORKSPACE); + + testAssertion.assertTest(null, projectId, API_KEY, TEST_WORKSPACE, List.of(), List.of(), List.of(), + List.of(), Map.of()); + } + + private Instant generateStartTime() { + return Instant.now().minusMillis(randomNumber(1, 1000)); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void findWithUsage(String endpoint, TracePageTestAssertion testAssertion) { + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .startTime(generateStartTime()) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(BigDecimal.ZERO) + .guardrailsValidations(null) + .build()) + .toList(); + traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); + + var traceIdToSpansMap = traces.stream() + .flatMap(trace -> PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(span -> span.toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .totalEstimatedCost(null) + .build())) + .collect(Collectors.groupingBy(Span::traceId)); + batchCreateSpansAndAssert( + traceIdToSpansMap.values().stream().flatMap(List::stream).toList(), API_KEY, TEST_WORKSPACE); + + traces = traces.stream().map(trace -> trace.toBuilder() + .usage(traceIdToSpansMap.get(trace.id()).stream() + .map(Span::usage) + .flatMap(usage -> usage.entrySet().stream()) + .collect(Collectors.groupingBy( + Map.Entry::getKey, Collectors.summingLong(Map.Entry::getValue)))) + .build()).toList(); + + var traceIdToCommentsMap = traces.stream() + .map(trace -> Pair.of(trace.id(), + IntStream.range(0, 5) + .mapToObj(i -> traceResourceClient.generateAndCreateComment(trace.id(), API_KEY, + TEST_WORKSPACE, 201)) + .toList())) + .collect(Collectors.toMap(Pair::getLeft, Pair::getRight)); + + traces = traces.stream().map(trace -> trace.toBuilder() + .usage(traceIdToSpansMap.get(trace.id()).stream() + .map(Span::usage) + .flatMap(usage -> usage.entrySet().stream()) + .collect(Collectors.groupingBy( + Map.Entry::getKey, Collectors.summingLong(Map.Entry::getValue)))) + .comments(traceIdToCommentsMap.get(trace.id())) + .build()).toList(); + + traces = updateSpanCounts(traces, traceIdToSpansMap); + + var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + + testAssertion.assertTest(projectName, null, API_KEY, TEST_WORKSPACE, values.expected(), values.unexpected(), + values.all(), List.of(), Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void findWithoutUsage(String endpoint, TracePageTestAssertion testAssertion) { + var apiKey = UUID.randomUUID().toString(); + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .startTime(generateStartTime()) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(BigDecimal.ZERO) + .guardrailsValidations(null) + .build()) + .toList(); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var spans = traces.stream() + .flatMap(trace -> PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(span -> span.toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .startTime(trace.startTime()) + .usage(null) + .totalEstimatedCost(null) + .build())) + .toList(); + batchCreateSpansAndAssert(spans, apiKey, workspaceName); + + traces = updateSpanCounts(traces, spans); + + var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), List.of(), Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + @DisplayName("when project name is not empty, then return traces by project name") + void whenProjectNameIsNotEmpty__thenReturnTracesByProjectName(String endpoint, + TracePageTestAssertion testAssertion) { + + var projectName = UUID.randomUUID().toString(); + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + List traces = new ArrayList<>(); + + for (int i = 0; i < 15; i++) { + Trace trace = createTrace() + .toBuilder() + .projectName(projectName) + .endTime(null) + .duration(null) + .output(null) + .tags(null) + .feedbackScores(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build(); + + traces.add(trace); + } + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), List.of(), Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + @DisplayName("when project id is not empty, then return traces by project id") + void whenProjectIdIsNotEmpty__thenReturnTracesByProjectId(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var projectName = UUID.randomUUID().toString(); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + Trace trace = createTrace() + .toBuilder() + .projectName(projectName) + .endTime(null) + .duration(null) + .output(null) + .projectId(null) + .tags(null) + .feedbackScores(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build(); + + create(trace, apiKey, workspaceName); + + UUID projectId = getProjectId(projectName, workspaceName, apiKey); + + var values = testAssertion.transformTestParams(List.of(), List.of(trace), List.of()); + + testAssertion.assertTest(null, projectId, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), List.of(), Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + @DisplayName("when filtering by workspace name, then return traces filtered") + void whenFilterWorkspaceName__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + + var workspaceName1 = UUID.randomUUID().toString(); + var workspaceName2 = UUID.randomUUID().toString(); + + var projectName1 = UUID.randomUUID().toString(); + + var workspaceId1 = UUID.randomUUID().toString(); + var workspaceId2 = UUID.randomUUID().toString(); + + var apiKey1 = UUID.randomUUID().toString(); + var apiKey2 = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey1, workspaceName1, workspaceId1); + mockTargetWorkspace(apiKey2, workspaceName2, workspaceId2); + + var traces1 = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName1) + .usage(null) + .threadId(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) + .comments(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .toList(); + + var traces2 = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName1) + .usage(null) + .threadId(null) + .feedbackScores(null) + .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .toList(); + + traceResourceClient.batchCreateTraces(traces1, apiKey1, workspaceName1); + traceResourceClient.batchCreateTraces(traces2, apiKey2, workspaceName2); + + var valueTraces1 = testAssertion.transformTestParams(traces1, traces1.reversed(), List.of()); + var valueTraces2 = testAssertion.transformTestParams(traces2, traces2.reversed(), List.of()); + + testAssertion.assertTest(projectName1, null, apiKey1, workspaceName1, valueTraces1.expected(), + valueTraces1.unexpected(), valueTraces1.all(), List.of(), Map.of()); + testAssertion.assertTest(projectName1, null, apiKey2, workspaceName2, valueTraces2.expected(), + valueTraces2.unexpected(), valueTraces2.all(), List.of(), Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + @DisplayName("when traces have cost estimation, then return total cost estimation") + void whenTracesHaveCostEstimation__thenReturnTotalCostEstimation(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var projectName = UUID.randomUUID().toString(); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + List traces = new ArrayList<>(); + + for (int i = 0; i < 5; i++) { + + Trace trace = createTrace() + .toBuilder() + .projectName(projectName) + .endTime(null) + .duration(null) + .output(null) + .projectId(null) + .tags(null) + .feedbackScores(null) + .usage(null) + .guardrailsValidations(null) + .totalEstimatedCost(BigDecimal.ZERO) + .build(); + + List spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(span -> span.toBuilder() + .usage(spanResourceClient.getTokenUsage()) + .model(spanResourceClient.randomModel().toString()) + .provider(spanResourceClient.provider()) + .traceId(trace.id()) + .projectName(projectName) + .feedbackScores(null) + .totalEstimatedCost(null) + .build()) + .toList(); + + batchCreateSpansAndAssert(spans, apiKey, workspaceName); + + Trace expectedTrace = trace.toBuilder() + .totalEstimatedCost(calculateEstimatedCost(spans)) + .usage(aggregateSpansUsage(spans)) + .build(); + + expectedTrace = updateSpanCounts(expectedTrace, spans); + + traces.add(expectedTrace); + } + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + UUID projectId = getProjectId(projectName, workspaceName, apiKey); + + var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + + testAssertion.assertTest(null, projectId, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), List.of(), Map.of()); + } + + @ParameterizedTest + @MethodSource("equalAndNotEqualFilters") + void whenFilterIdAndNameEqual__thenReturnTracesFiltered(String endpoint, + Operator operator, + Function, List> getExpectedTraces, + Function, List> getUnexpectedTraces, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = getExpectedTraces.apply(traces); + var unexpectedTraces = getUnexpectedTraces.apply(traces); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.ID) + .operator(operator) + .value(traces.getFirst().id().toString()) + .build(), + TraceFilter.builder() + .field(TraceField.NAME) + .operator(operator) + .value(traces.getFirst().name()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("equalAndNotEqualFilters") + void whenFilterByThreadEqual__thenReturnTracesFiltered(String endpoint, + Operator operator, + Function, List> getExpectedTraces, + Function, List> getUnexpectedTraces, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(UUID.randomUUID().toString()) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(traces.size() - 1, traces.getLast().toBuilder() + .threadId(null) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = getExpectedTraces.apply(traces); + var unexpectedTraces = getUnexpectedTraces.apply(traces); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.THREAD_ID) + .operator(operator) + .value(traces.getFirst().threadId()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterNameEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .totalEstimatedCost(null) + .feedbackScores(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + List filters = List.of(TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.EQUAL) + .value(traces.getFirst().name().toUpperCase()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterNameStartsWith__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.STARTS_WITH) + .value(traces.getFirst().name().substring(0, traces.getFirst().name().length() - 4).toUpperCase()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterNameEndsWith__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.ENDS_WITH) + .value(traces.getFirst().name().substring(3).toUpperCase()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterNameContains__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.CONTAINS) + .value(traces.getFirst().name().substring(2, traces.getFirst().name().length() - 3).toUpperCase()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterAnnotationQueueIdContains__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var project = factory.manufacturePojo(Project.class); + var projectId = projectResourceClient.createProject(project, apiKey, workspaceName); + + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(project.name()) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .projectName(project.name()) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + // Create annotation queue with items + var queue1 = prepareAnnotationQueue(projectId); + var queue2 = prepareAnnotationQueue(projectId); + annotationQueuesResourceClient.createAnnotationQueueBatch( + new LinkedHashSet<>(List.of(queue1, queue2)), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); + + annotationQueuesResourceClient.addItemsToAnnotationQueue( + queue1.id(), Set.of(traces.getFirst().id()), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); + annotationQueuesResourceClient.addItemsToAnnotationQueue( + queue2.id(), Set.of(traces.get(1).id()), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.ANNOTATION_QUEUE_IDS) + .operator(Operator.CONTAINS) + .value(queue1.id().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(project.name(), null, apiKey, workspaceName, values.expected(), + values.unexpected(), + values.all(), + filters, Map.of()); + } + + private AnnotationQueue prepareAnnotationQueue(UUID projectId) { + return factory.manufacturePojo(AnnotationQueue.class) + .toBuilder() + .projectId(projectId) + .scope(AnnotationQueue.AnnotationScope.TRACE) + .build(); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterNameNotContains__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traceName = generator.generate().toString(); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .name(traceName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .name(generator.generate().toString()) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.NOT_CONTAINS) + .value(traceName.toUpperCase()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("equalAndNotEqualFilters") + void whenFilterStartTimeEqual__thenReturnTracesFiltered(String endpoint, + Operator operator, + Function, List> getExpectedTraces, + Function, List> getUnexpectedTraces, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = getExpectedTraces.apply(traces); + var unexpectedTraces = getUnexpectedTraces.apply(traces); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.START_TIME) + .operator(operator) + .value(traces.getFirst().startTime().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterStartTimeGreaterThan__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .startTime(Instant.now().minusSeconds(60 * 5)) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .startTime(Instant.now().plusSeconds(60 * 5)) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.START_TIME) + .operator(Operator.GREATER_THAN) + .value(Instant.now().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterStartTimeGreaterThanEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .startTime(Instant.now().minusSeconds(60 * 5)) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .startTime(Instant.now()) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.START_TIME) + .operator(Operator.GREATER_THAN_EQUAL) + .value(traces.getFirst().startTime().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterStartTimeLessThan__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .startTime(Instant.now().plusSeconds(60 * 5)) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .startTime(Instant.now().minusSeconds(60 * 5)) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.START_TIME) + .operator(Operator.LESS_THAN) + .value(Instant.now().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterStartTimeLessThanEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .startTime(Instant.now().plusSeconds(60 * 5)) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .startTime(Instant.now().minusSeconds(60 * 5)) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.START_TIME) + .operator(Operator.LESS_THAN_EQUAL) + .value(traces.getFirst().startTime().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterEndTimeEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.END_TIME) + .operator(Operator.EQUAL) + .value(traces.getFirst().endTime().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterInputEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .totalEstimatedCost(null) + .feedbackScores(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.INPUT) + .operator(Operator.EQUAL) + .value(traces.getFirst().input().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterOutputEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .totalEstimatedCost(null) + .feedbackScores(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.OUTPUT) + .operator(Operator.EQUAL) + .value(traces.getFirst().output().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterTotalEstimatedCostGreaterThen__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(spanInStream -> spanInStream.toBuilder() + .projectName(projectName) + .traceId(traces.getFirst().id()) + .usage(Map.of("completion_tokens", Math.abs(factory.manufacturePojo(Integer.class)), + "prompt_tokens", Math.abs(factory.manufacturePojo(Integer.class)))) + .model("gpt-3.5-turbo-1106") + .provider("openai") + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toList()); + + batchCreateSpansAndAssert(spans, apiKey, workspaceName); + + var finalTraces = updateSpanCounts(traces, spans); + var unexpectedTraces = finalTraces.subList(1, traces.size()); + var expectedTrace = finalTraces.getFirst().toBuilder() + .usage(aggregateSpansUsage(spans)) + .totalEstimatedCost(calculateEstimatedCost(spans)) + .build(); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.TOTAL_ESTIMATED_COST) + .operator(Operator.GREATER_THAN) + .value("0") + .build()); + + var values = testAssertion.transformTestParams(finalTraces, List.of(expectedTrace), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("equalAndNotEqualFilters") + void whenFilterTotalEstimatedCostEqual_NotEqual__thenReturnTracesFiltered(String endpoint, + Operator operator, + Function, List> getUnexpectedTraces, // Here we swap the expected and unexpected traces + Function, List> getExpectedTraces, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(spanInStream -> spanInStream.toBuilder() + .projectName(projectName) + .traceId(traces.getFirst().id()) + .usage(Map.of("completion_tokens", Math.abs(factory.manufacturePojo(Integer.class)), + "prompt_tokens", Math.abs(factory.manufacturePojo(Integer.class)))) + .model("gpt-3.5-turbo-1106") + .provider("openai") + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toList()); + + var otherSpans = traces.stream().skip(1) + .flatMap(trace -> PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(span -> span.toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .usage(null) + .model(null) + .totalEstimatedCost(null) + .build())) + .toList(); + + var allSpans = Stream.concat(spans.stream(), otherSpans.stream()).toList(); + batchCreateSpansAndAssert(allSpans, apiKey, workspaceName); + + traces.set(0, traces.getFirst().toBuilder() + .usage(aggregateSpansUsage(spans)) + .totalEstimatedCost(calculateEstimatedCost(spans)) + .build()); + + var finalTraces = updateSpanCounts(traces, allSpans); + var expectedTraces = getExpectedTraces.apply(finalTraces); + var unexpectedTraces = getUnexpectedTraces.apply(finalTraces); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.TOTAL_ESTIMATED_COST) + .operator(operator) + .value("0.00") + .build()); + + var values = testAssertion.transformTestParams(finalTraces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + Stream whenFilterLlmSpanCountOperator__thenReturnTracesFiltered() { + return getFilterTestArguments().flatMap(args -> Stream.of( + Arguments.of(args.get()[0], args.get()[1], Operator.EQUAL), + Arguments.of(args.get()[0], args.get()[1], Operator.NOT_EQUAL), + Arguments.of(args.get()[0], args.get()[1], Operator.GREATER_THAN), + Arguments.of(args.get()[0], args.get()[1], Operator.GREATER_THAN_EQUAL), + Arguments.of(args.get()[0], args.get()[1], Operator.LESS_THAN), + Arguments.of(args.get()[0], args.get()[1], Operator.LESS_THAN_EQUAL))); + } + + @ParameterizedTest + @MethodSource + void whenFilterLlmSpanCountOperator__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + Operator operator) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> { + var llmSpanCount = RandomUtils.secure().randomInt(1, 7); + return trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .spanCount(llmSpanCount + RandomUtils.secure().randomInt(1, 7)) + .llmSpanCount(llmSpanCount) + .build(); + }) + .toList(); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var spans = traces.stream() + .flatMap(trace -> IntStream.range(0, trace.spanCount()) + .mapToObj(i -> factory.manufacturePojo(Span.class).toBuilder() + .usage(null) + .totalEstimatedCost(null) + .projectName(projectName) + .traceId(trace.id()) + .type(i < trace.llmSpanCount() ? SpanType.llm : SpanType.general) + .build())) + .toList(); + + spanResourceClient.batchCreateSpans(spans, apiKey, workspaceName); + + var llmSpanCountToCompareAgainst = traces.getFirst().llmSpanCount(); + + Predicate matchesFilter = makeLlmSpanCountPredicate(operator, llmSpanCountToCompareAgainst); + Comparator traceIdComparator = Comparator.comparing(Trace::id).reversed(); + + var expectedTraces = traces.stream() + .filter(matchesFilter) + .sorted(traceIdComparator) + .collect(Collectors.toList()); + + var unexpectedTraces = traces.stream() + .filter(matchesFilter.negate()) + .sorted(traceIdComparator) + .collect(Collectors.toList()); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.LLM_SPAN_COUNT) + .operator(operator) + .value(Integer.toString(llmSpanCountToCompareAgainst)) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + Predicate makeLlmSpanCountPredicate(Operator operator, int value) { + switch (operator) { + case Operator.EQUAL : + return trace -> trace.llmSpanCount() == value; + case Operator.NOT_EQUAL : + return trace -> trace.llmSpanCount() != value; + case Operator.GREATER_THAN : + return trace -> trace.llmSpanCount() > value; + case Operator.GREATER_THAN_EQUAL : + return trace -> trace.llmSpanCount() >= value; + case Operator.LESS_THAN : + return trace -> trace.llmSpanCount() < value; + case Operator.LESS_THAN_EQUAL : + return trace -> trace.llmSpanCount() <= value; + default : + throw new IllegalArgumentException("Invalid operator for llm span count filtering: " + operator); + } + } + + @ParameterizedTest + @MethodSource("equalAndNotEqualFilters") + void whenFilterMetadataEqualString__thenReturnTracesFiltered(String endpoint, + Operator operator, + Function, List> getExpectedTraces, + Function, List> getUnexpectedTraces, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + traces.forEach(trace -> create(trace, apiKey, workspaceName)); + var expectedTraces = getExpectedTraces.apply(traces); + var unexpectedTraces = getUnexpectedTraces.apply(traces); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(operator) + .key("$.model[0].version") + .value("OPENAI, CHAT-GPT 4.0") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataEqualNumber__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2023,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.EQUAL) + .key("model[0].year") + .value("2023") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataEqualBoolean__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata( + JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":false,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.EQUAL) + .key("model[0].year") + .value("TRUE") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataEqualNull__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .threadId(null) + .totalEstimatedCost(null) + .feedbackScores(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.EQUAL) + .key("model[0].year") + .value("NULL") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataContainsString__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.CONTAINS) + .key("model[0].version") + .value("CHAT-GPT") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataContainsNumber__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .threadId(null) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":\"two thousand twenty " + + "four\",\"version\":\"OpenAI, Chat-GPT 4.0\"}]}")) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2023,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.CONTAINS) + .key("model[0].year") + .value("02") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataContainsBoolean__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata( + JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":false,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.CONTAINS) + .key("model[0].year") + .value("TRU") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataContainsNull__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + + "version\"}]}")) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.CONTAINS) + .key("model[0].year") + .value("NUL") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataGreaterThanNumber__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2020," + + "\"version\":\"OpenAI, Chat-GPT 4.0\"}]}")) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.GREATER_THAN) + .key("model[0].year") + .value("2023") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataGreaterThanString__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"openAI, " + + "Chat-GPT 4.0\"}]}")) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.GREATER_THAN) + .key("model[0].version") + .value("a") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataGreaterThanBoolean__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"openAI, " + + "Chat-GPT 4.0\"}]}")) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.GREATER_THAN) + .key("model[0].year") + .value("a") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataGreaterThanNull__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"openAI, " + + "Chat-GPT 4.0\"}]}")) + .feedbackScores(null) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.GREATER_THAN) + .key("model[0].year") + .value("a") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataLessThanNumber__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2026," + + "\"version\":\"OpenAI, Chat-GPT 4.0\"}]}")) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(0, traces.getFirst().toBuilder() + .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\"}]}")) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.LESS_THAN) + .key("model[0].year") + .value("2025") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataLessThanString__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"openAI, " + + "Chat-GPT 4.0\"}]}")) + .feedbackScores(null) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.LESS_THAN) + .key("model[0].version") + .value("z") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataLessThanBoolean__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"openAI, " + + "Chat-GPT 4.0\"}]}")) + .feedbackScores(null) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.LESS_THAN) + .key("model[0].year") + .value("z") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterMetadataLessThanNull__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .metadata(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"openAI, " + + "Chat-GPT 4.0\"}]}")) + .feedbackScores(null) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.METADATA) + .operator(Operator.LESS_THAN) + .key("model[0].year") + .value("z") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterTagsContains__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.TAGS) + .operator(Operator.CONTAINS) + .value(traces.getFirst().tags().stream() + .toList() + .get(2) + .substring(0, traces.getFirst().name().length() - 4) + .toUpperCase()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getUsageKeyArgs") + void whenFilterUsageEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + String usageKey, + Field field) { + + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var otherUsageValue = randomNumber(1, 8); + var usageValue = randomNumber(); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(Map.of(usageKey, (long) otherUsageValue)) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toList()); + + traces.set(0, traces.getFirst().toBuilder() + .usage(Map.of(usageKey, (long) usageValue)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var traceIdToSpanMap = traces.stream() + .map(trace -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .usage(Map.of(usageKey, otherUsageValue)) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toMap(Span::traceId, Function.identity())); + traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() + .usage(Map.of(usageKey, usageValue)) + .build()); + batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); + + traces = updateSpanCounts(traces, traceIdToSpanMap.values().stream().toList()); + var expectedTraces = List.of(traces.getFirst()); + var unrelatedTraces = List.of(createTrace()); + + traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(field) + .operator(Operator.EQUAL) + .value(traces.getFirst().usage().get(usageKey).toString()) + .build()); + + var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) + .toList(); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getUsageKeyArgs") + void whenFilterUsageGreaterThan__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + String usageKey, + Field field) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(Map.of(usageKey, 123L)) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(1) + .spanCount(1) + .build()) + .collect(Collectors.toList()); + traces.set(0, traces.getFirst().toBuilder() + .usage(Map.of(usageKey, 456L)) + .build()); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var traceIdToSpanMap = traces.stream() + .map(trace -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .usage(Map.of(usageKey, 123)) + .totalEstimatedCost(null) + .type(SpanType.llm) + .build()) + .collect(Collectors.toMap(Span::traceId, Function.identity())); + traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() + .usage(Map.of(usageKey, 456)) + .build()); + batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + var unrelatedTraces = List.of(createTrace()); + + traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(field) + .operator(Operator.GREATER_THAN) + .value("123") + .build()); + + var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) + .toList(); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getUsageKeyArgs") + void whenFilterUsageGreaterThanEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + String usageKey, + Field field) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(Map.of(usageKey, 123L)) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toList()); + traces.set(0, traces.getFirst().toBuilder() + .usage(Map.of(usageKey, 456L)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var traceIdToSpanMap = traces.stream() + .map(trace -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .usage(Map.of(usageKey, 123)) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toMap(Span::traceId, Function.identity())); + traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() + .usage(Map.of(usageKey, 456)) + .build()); + batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); + + traces = updateSpanCounts(traces, traceIdToSpanMap.values().stream().toList()); + var expectedTraces = List.of(traces.getFirst()); + var unrelatedTraces = List.of(createTrace()); + + traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(field) + .operator(Operator.GREATER_THAN_EQUAL) + .value(traces.getFirst().usage().get(usageKey).toString()) + .build()); + + var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) + .toList(); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getUsageKeyArgs") + void whenFilterUsageLessThan__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + String usageKey, + Field field) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(Map.of(usageKey, 456L)) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toList()); + traces.set(0, traces.getFirst().toBuilder() + .usage(Map.of(usageKey, 123L)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var traceIdToSpanMap = traces.stream() + .map(trace -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .usage(Map.of(usageKey, 456)) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toMap(Span::traceId, Function.identity())); + traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() + .usage(Map.of(usageKey, 123)) + .build()); + batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); + + traces = updateSpanCounts(traces, traceIdToSpanMap.values().stream().toList()); + var expectedTraces = List.of(traces.getFirst()); + var unrelatedTraces = List.of(createTrace()); + + traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(field) + .operator(Operator.LESS_THAN) + .value("456") + .build()); + + var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) + .toList(); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getUsageKeyArgs") + void whenFilterUsageLessThanEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + String usageKey, + Field field) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(Map.of(usageKey, 456L)) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(1) + .spanCount(1) + .build()) + .collect(Collectors.toList()); + traces.set(0, traces.getFirst().toBuilder() + .usage(Map.of(usageKey, 123L)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var traceIdToSpanMap = traces.stream() + .map(trace -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(projectName) + .traceId(trace.id()) + .usage(Map.of(usageKey, 456)) + .totalEstimatedCost(null) + .type(SpanType.llm) + .build()) + .collect(Collectors.toMap(Span::traceId, Function.identity())); + traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() + .usage(Map.of(usageKey, 123)) + .build()); + batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + var unrelatedTraces = List.of(createTrace()); + + traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(field) + .operator(Operator.LESS_THAN_EQUAL) + .value(traces.getFirst().usage().get(usageKey).toString()) + .build()); + + var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) + .toList(); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFeedbackScoresArgs") + void whenFilterFeedbackScoresEqual__thenReturnTracesFiltered(String endpoint, + Operator operator, + Function, List> getExpectedTraces, + Function, List> getUnexpectedTraces, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .totalEstimatedCost(null) + .feedbackScores(trace.feedbackScores().stream() + .map(feedbackScore -> feedbackScore.toBuilder() + .value(factory.manufacturePojo(BigDecimal.class)) + .build()) + .collect(Collectors.toList())) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(1, traces.get(1).toBuilder() + .feedbackScores( + updateFeedbackScore(traces.get(1).feedbackScores(), traces.getFirst().feedbackScores(), 2)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + traces.forEach(trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var expectedTraces = getExpectedTraces.apply(traces); + var unexpectedTraces = getUnexpectedTraces.apply(traces); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(operator) + .key(traces.getFirst().feedbackScores().get(1).name().toUpperCase()) + .value(traces.getFirst().feedbackScores().get(1).value().toString()) + .build(), + TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(operator) + .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) + .value(traces.getFirst().feedbackScores().get(2).value().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource + void getTracesByProject__whenFilterFeedbackScoresIsEmpty__thenReturnTracesFiltered( + Operator operator, + Function, List> getExpectedTraces, + Function, List> getUnexpectedTraces, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .feedbackScores(trace.feedbackScores().stream() + .map(feedbackScore -> feedbackScore.toBuilder() + .value(factory.manufacturePojo(BigDecimal.class)) + .build()) + .collect(Collectors.toList())) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traces.set(traces.size() - 1, traces.getLast().toBuilder().feedbackScores(null).build()); + traces.forEach(trace1 -> create(trace1, apiKey, workspaceName)); + traces.subList(0, traces.size() - 1).forEach(trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + var expectedTraces = getExpectedTraces.apply(traces); + var unexpectedTraces = getUnexpectedTraces.apply(traces); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(operator) + .key(traces.getFirst().feedbackScores().getFirst().name()) + .value("") + .build()); + var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + private Stream getTracesByProject__whenFilterFeedbackScoresIsEmpty__thenReturnTracesFiltered() { + return Stream.of( + Arguments.of(Operator.IS_NOT_EMPTY, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceTestAssertion), + Arguments.of(Operator.IS_EMPTY, + (Function, List>) traces -> traces.subList(1, traces.size()), + (Function, List>) traces -> List.of(traces.getFirst()), + traceTestAssertion), + Arguments.of(Operator.IS_NOT_EMPTY, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceStatsAssertion), + Arguments.of(Operator.IS_EMPTY, + (Function, List>) traces -> traces.subList(1, traces.size()), + (Function, List>) traces -> List.of(traces.getFirst()), + traceStatsAssertion), + Arguments.of(Operator.IS_NOT_EMPTY, + (Function, List>) traces -> List.of(traces.getFirst()), + (Function, List>) traces -> traces.subList(1, traces.size()), + traceStreamTestAssertion), + Arguments.of(Operator.IS_EMPTY, + (Function, List>) traces -> traces.subList(1, traces.size()), + (Function, List>) traces -> List.of(traces.getFirst()), + traceStreamTestAssertion)); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterFeedbackScoresGreaterThan__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .totalEstimatedCost(null) + .llmSpanCount(0) + .spanCount(0) + .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() + .map(feedbackScore -> feedbackScore.toBuilder() + .value(factory.manufacturePojo(BigDecimal.class)) + .build()) + .collect(Collectors.toList()), 2, 1234.5678)) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 2345.6789)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + traces.forEach(trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + unexpectedTraces.forEach( + trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.EQUAL) + .value(traces.getFirst().name()) + .build(), + TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(Operator.GREATER_THAN) + .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) + .value("2345.6788") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterFeedbackScoresGreaterThanEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() + .map(feedbackScore -> feedbackScore.toBuilder() + .value(factory.manufacturePojo(BigDecimal.class)) + .build()) + .collect(Collectors.toList()), 2, 1234.5678)) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 2345.6789)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + traces.forEach(trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + unexpectedTraces.forEach( + trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(Operator.GREATER_THAN_EQUAL) + .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) + .value(traces.getFirst().feedbackScores().get(2).value().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterFeedbackScoresLessThan__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .comments(null) + .totalEstimatedCost(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() + .map(feedbackScore -> feedbackScore.toBuilder() + .value(factory.manufacturePojo(BigDecimal.class)) + .build()) + .collect(Collectors.toList()), 2, 2345.6789)) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 1234.5678)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + traces.forEach(trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + unexpectedTraces.forEach( + trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(Operator.LESS_THAN) + .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) + .value("2345.6788") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterFeedbackScoresLessThanEqual__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .totalEstimatedCost(null) + .llmSpanCount(0) + .spanCount(0) + .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() + .map(feedbackScore -> feedbackScore.toBuilder() + .value(factory.manufacturePojo(BigDecimal.class)) + .build()) + .collect(Collectors.toList()), 2, 2345.6789)) + .guardrailsValidations(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 1234.5678)) + .build());; + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + traces.forEach(trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectName(RandomStringUtils.secure().nextAlphanumeric(20)) + .projectId(null) + .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) + .build()); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + unexpectedTraces.forEach( + trace -> trace.feedbackScores() + .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.FEEDBACK_SCORES) + .operator(Operator.LESS_THAN_EQUAL) + .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) + .value(traces.getFirst().feedbackScores().get(2).value().toString()) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getDurationArgs") + void whenFilterByDuration__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion, + Operator operator, + long end, + double duration) { + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = generator.generate().toString(); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> { + Instant now = Instant.now(); + return trace.toBuilder() + .projectId(null) + .usage(null) + .projectName(projectName) + .feedbackScores(null) + .threadId(null) + .totalEstimatedCost(null) + .startTime(now) + .endTime(Set.of(Operator.LESS_THAN, Operator.LESS_THAN_EQUAL).contains(operator) + ? Instant.now().plusSeconds(2) + : now.plusNanos(1000)) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build(); + }) + .collect(Collectors.toCollection(ArrayList::new)); + + var start = Instant.now().truncatedTo(ChronoUnit.MILLIS); + traces.set(0, traces.getFirst().toBuilder() + .startTime(start) + .endTime(start.plus(end, ChronoUnit.MICROS)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = List.of(traces.getFirst()); + + var unexpectedTraces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(span -> span.toBuilder() + .projectId(null) + .build()) + .toList(); + + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of( + TraceFilter.builder() + .field(TraceField.DURATION) + .operator(operator) + .value(String.valueOf(duration)) + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterInvalidOperatorForFieldTypeArgs") + void whenFilterInvalidOperatorForFieldType__thenReturn400(String path, TraceFilter filter) { + + String errorMessage = filter.field().getType() == FieldType.CUSTOM + ? "Invalid key '%s' for custom filter".formatted(filter.key()) + : "Invalid operator '%s' for field '%s' of type '%s'".formatted( + filter.operator().getQueryParamOperator(), + filter.field().getQueryParamField(), + filter.field().getType()); + + var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( + HttpStatus.SC_BAD_REQUEST, errorMessage); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var filters = List.of(filter); + + Response actualResponse; + if (path.equals("/search")) { + actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path(path) + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .post(Entity.json(TraceSearchStreamRequest.builder() + .projectName(projectName) + .filters(filters) + .build())); + + } else { + + actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path(path) + .queryParam("project_name", projectName) + .queryParam("filters", toURLEncodedQueryParam(filters)) + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .get(); + } + + try (actualResponse) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + + var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(actualError).isEqualTo(expectedError); + } + + } + + @ParameterizedTest + @MethodSource("getFilterInvalidValueOrKeyForFieldTypeArgs") + void whenFilterInvalidValueOrKeyForFieldType__thenReturn400(String path, TraceFilter filter) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( + 400, + "Invalid value '%s' or key '%s' for field '%s' of type '%s'".formatted( + filter.value(), + filter.key(), + filter.field().getQueryParamField(), + filter.field().getType())); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var filters = List.of(filter); + + Response actualResponse; + + if (path.equals("/search")) { + + actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path(path) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(TraceSearchStreamRequest.builder() + .projectName(projectName) + .filters(filters) + .build())); + + } else { + actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path(path) + .queryParam("project_name", projectName) + .queryParam("filters", toURLEncodedQueryParam(filters)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get(); + } + + try (actualResponse) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + + var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(actualError).isEqualTo(expectedError); + } + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterGuardrails__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .threadId(null) + .totalEstimatedCost(null) + .feedbackScores(null) + .guardrailsValidations(null) + .comments(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var guardrailsByTraceId = traces.stream() + .collect(Collectors.toMap(Trace::id, trace -> guardrailsGenerator.generateGuardrailsForTrace( + trace.id(), randomUUID(), trace.projectName()))); + + // set the first trace with failed guardrails + guardrailsByTraceId.put(traces.getFirst().id(), guardrailsByTraceId.get(traces.getFirst().id()).stream() + .map(guardrail -> guardrail.toBuilder().result(GuardrailResult.FAILED).build()) + .toList()); + + // set the rest of traces with passed guardrails + traces.subList(1, traces.size()).forEach(trace -> guardrailsByTraceId.put(trace.id(), + guardrailsByTraceId.get(trace.id()).stream() + .map(guardrail -> guardrail.toBuilder() + .result(GuardrailResult.PASSED) + .build()) + .toList())); + + guardrailsByTraceId.values() + .forEach(guardrail -> guardrailsResourceClient.addBatch(guardrail, apiKey, + workspaceName)); + + traces = traces.stream().map(trace -> trace.toBuilder() + .guardrailsValidations(GuardrailsMapper.INSTANCE.mapToValidations( + guardrailsByTraceId.get(trace.id()))) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + // assert failed guardrails + var filtersFailed = List.of( + TraceFilter.builder() + .field(TraceField.GUARDRAILS) + .operator(Operator.EQUAL) + .value(GuardrailResult.FAILED.getResult()) + .build()); + + var valuesFailed = testAssertion.transformTestParams(traces, List.of(traces.getFirst()), + traces.subList(1, traces.size())); + testAssertion.assertTest(projectName, null, apiKey, workspaceName, valuesFailed.expected(), + valuesFailed.unexpected(), valuesFailed.all(), filtersFailed, Map.of()); + + // assert passed guardrails + var filtersPassed = List.of( + TraceFilter.builder() + .field(TraceField.GUARDRAILS) + .operator(Operator.EQUAL) + .value(GuardrailResult.PASSED.getResult()) + .build()); + + var valuesPassed = testAssertion.transformTestParams(traces, traces.subList(1, traces.size()).reversed(), + List.of(traces.getFirst())); + testAssertion.assertTest(projectName, null, apiKey, workspaceName, valuesPassed.expected(), + valuesPassed.unexpected(), valuesPassed.all(), filtersPassed, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterErrorIsNotEmpty__thenReturnTracesFiltered(String endpoint, + TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .errorInfo(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .errorInfo(factory.manufacturePojo(ErrorInfo.class)) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.ERROR_INFO) + .operator(Operator.IS_NOT_EMPTY) + .value("") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + + @ParameterizedTest + @MethodSource("getFilterTestArguments") + void whenFilterErrorIsEmpty__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .totalEstimatedCost(null) + .threadId(null) + .guardrailsValidations(null) + .llmSpanCount(0) + .spanCount(0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .errorInfo(null) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var expectedTraces = List.of(traces.getFirst()); + var unexpectedTraces = List.of(createTrace().toBuilder() + .projectId(null) + .build()); + traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + + var filters = List.of(TraceFilter.builder() + .field(TraceField.ERROR_INFO) + .operator(Operator.IS_EMPTY) + .value("") + .build()); + + var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), + values.all(), + filters, Map.of()); + } + } + + private BigDecimal calculateEstimatedCost(List spans) { + return spans.stream() + .map(span -> CostService.calculateCost(span.model(), span.provider(), span.usage(), null)) + .reduce(BigDecimal.ZERO, BigDecimal::add); + } + + private String getValidValue(Field field) { + return switch (field.getType()) { + case STRING, LIST, DICTIONARY, CUSTOM, ENUM, STRING_STATE_DB -> + RandomStringUtils.secure().nextAlphanumeric(10); + case NUMBER, DURATION, FEEDBACK_SCORES_NUMBER -> String.valueOf(randomNumber(1, 10)); + case DATE_TIME, DATE_TIME_STATE_DB -> Instant.now().toString(); + case ERROR_CONTAINER -> ""; + }; + } + + private String getKey(Field field) { + return switch (field.getType()) { + case STRING, NUMBER, DURATION, DATE_TIME, LIST, ENUM, ERROR_CONTAINER, STRING_STATE_DB, DATE_TIME_STATE_DB, + DICTIONARY -> + null; + case FEEDBACK_SCORES_NUMBER, CUSTOM -> RandomStringUtils.secure().nextAlphanumeric(10); + }; + } + + private String getInvalidValue(Field field) { + return switch (field.getType()) { + case STRING, DICTIONARY, CUSTOM, LIST, ENUM, ERROR_CONTAINER, STRING_STATE_DB, DATE_TIME_STATE_DB -> " "; + case NUMBER, DURATION, DATE_TIME, FEEDBACK_SCORES_NUMBER -> RandomStringUtils.secure().nextAlphanumeric(10); + }; + } + + @Nested + @DisplayName("Find traces:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class FindTraces { + + @ParameterizedTest + @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments") + void findWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) { + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = Stream.of(createTrace()) + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(null) + .input(original) + .output(original) + .metadata(original) + .build()) + .toList(); + + traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); + + var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .queryParam("page", 1) + .queryParam("size", 5) + .queryParam("project_name", projectName) + .queryParam("truncate", truncate) + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .get(); + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + + var actualPage = actualResponse.readEntity(Trace.TracePage.class); + var actualTraces = actualPage.content(); + + assertThat(actualTraces).hasSize(1); + + var expectedTraces = traces.stream() + .map(trace -> trace.toBuilder() + .input(expected) + .output(expected) + .metadata(expected) + .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(trace.startTime(), + trace.endTime())) + .build()) + .toList(); + + assertThat(actualTraces) + .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS_TRACES) + .containsExactlyElementsOf(expectedTraces); + } + + @ParameterizedTest + @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments") + void searchWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) { + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = Stream.of(createTrace()) + .map(trace -> trace.toBuilder() + .projectName(projectName) + .usage(null) + .input(original) + .output(original) + .metadata(original) + .build()) + .toList(); + + traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); + + TraceSearchStreamRequest streamRequest = TraceSearchStreamRequest.builder() + .truncate(truncate) + .projectName(projectName) + .limit(5) + .build(); + + var actualTraces = traceResourceClient.getStreamAndAssertContent(API_KEY, TEST_WORKSPACE, streamRequest); + + assertThat(actualTraces).hasSize(1); + + var expectedTraces = traces.stream() + .map(trace -> trace.toBuilder() + .input(expected) + .output(expected) + .metadata(expected) + .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(trace.startTime(), + trace.endTime())) + .build()) + .toList(); + + TraceAssertions.assertTraces(actualTraces, expectedTraces, USER); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void whenUsingPagination__thenReturnTracesPaginated(boolean stream) { + + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .comments(null) + .totalEstimatedCost(null) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var expectedTraces = traces.stream() + .sorted(Comparator.comparing(Trace::id).reversed()) + .toList(); + + int pageSize = 2; + + if (stream) { + AtomicReference lastId = new AtomicReference<>(null); + Lists.partition(expectedTraces, pageSize) + .forEach(trace -> { + var actualTraces = traceResourceClient.getStreamAndAssertContent(apiKey, workspaceName, + TraceSearchStreamRequest.builder() + .projectName(projectName) + .lastRetrievedId(lastId.get()) + .limit(pageSize) + .build()); + + TraceAssertions.assertTraces(actualTraces, trace, USER); + + lastId.set(actualTraces.getLast().id()); + }); + } else { + for (int i = 0; i < expectedTraces.size() / pageSize; i++) { + int page = i + 1; + getAndAssertPage( + page, + pageSize, + projectName, + null, + List.of(), + expectedTraces.subList(i * pageSize, Math.min((i + 1) * pageSize, expectedTraces.size())), + List.of(), + workspaceName, + apiKey, + List.of(), + traces.size(), Set.of()); + } + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void whenFilterByVisibilityScoreEqual__thenReturnTracesFiltered(boolean stream) { + + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .comments(null) + .totalEstimatedCost(null) + .build()) + .toList(); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + TraceFilter filter = TraceFilter.builder() + .field(TraceField.VISIBILITY_MODE) + .operator(Operator.EQUAL) + .value(VisibilityMode.DEFAULT.getValue()) + .build(); + + var actualTraces = traceResourceClient.getStreamAndAssertContent(apiKey, workspaceName, + TraceSearchStreamRequest.builder() + .projectName(projectName) + .filters(List.of(filter)) + .build()); + + if (stream) { + TraceAssertions.assertTraces(actualTraces, traces.reversed(), USER); + } else { + getAndAssertPage( + 1, + 100, + projectName, + null, + List.of(filter), + traces.reversed(), + List.of(), + workspaceName, + apiKey, + List.of(), + traces.size(), Set.of()); + } + } + + @ParameterizedTest + @MethodSource + void whenFilterByCustomFilter__thenReturnTracesFiltered(String key, String value, Operator operator) { + + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .comments(null) + .totalEstimatedCost(null) + .build()) + .collect(toCollection(ArrayList::new)); + + traces.set(0, traces.getFirst().toBuilder() + .input(JsonUtils + .getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + + "Chat-GPT 4.0\",\"trueFlag\":true,\"nullField\":null}]}")) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + TraceFilter filter = TraceFilter.builder() + .field(CUSTOM) + .operator(operator) + .key(key) + .value(value) + .build(); + + getAndAssertPage( + 1, + 100, + projectName, + null, + List.of(filter), + List.of(traces.getFirst()), + traces.subList(1, traces.size()), + workspaceName, + apiKey, + List.of(), + 1, Set.of()); + } + + private Stream whenFilterByCustomFilter__thenReturnTracesFiltered() { + return Stream.of( + Arguments.of( + "input.model[0].year", + "2024", + Operator.EQUAL), + Arguments.of( + "input.model[0].year", + "2025", + Operator.LESS_THAN), + Arguments.of( + "input", + "Chat-GPT 4.0", + Operator.CONTAINS)); + } + + @ParameterizedTest + @MethodSource + void getTracesByProject__whenSortingByValidFields__thenReturnTracesSorted(Comparator comparator, + SortingField sorting) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> { + var llmSpanCount = RandomUtils.secure().randomInt(1, 7); + return trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) + .comments(null) + .spanCount(llmSpanCount + RandomUtils.secure().randomInt(1, 7)) + .llmSpanCount(llmSpanCount) + .build(); + }) + .map(trace -> trace.toBuilder() + .duration(trace.startTime().until(trace.endTime(), ChronoUnit.MICROS) / 1000.0) + .build()) + .toList(); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + var spans = traces.stream() + .flatMap(trace -> IntStream.range(0, trace.spanCount()) + .mapToObj(i -> factory.manufacturePojo(Span.class).toBuilder() + .usage(Map.of("completion_tokens", RandomUtils.secure().randomInt())) + .projectName(projectName) + .traceId(trace.id()) + .type(i < trace.llmSpanCount() ? SpanType.llm : SpanType.general) + .build())) + .toList(); + + spanResourceClient.batchCreateSpans(spans, apiKey, workspaceName); + + var spansByTrace = spans.stream().collect(Collectors.groupingBy(Span::traceId)); + traces = traces.stream() + .map(t -> t.toBuilder() + .usage(aggregateSpansUsage(spansByTrace.get(t.id()))) + .build()) + .toList(); + + var expectedTraces = traces.stream() + .sorted(comparator) + .toList(); + + List sortingFields = List.of(sorting); + + getAndAssertPage(workspaceName, projectName, null, List.of(), traces, expectedTraces, List.of(), apiKey, + sortingFields, Set.of()); + } + + @Test + void createAndRetrieveTraces__spanCountReflectsActualSpans_andTotalCountMatches() { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + + // Create traces with varying spanCount values + List traces = IntStream.range(0, 5) + .mapToObj(i -> createTrace().toBuilder() + .projectId(null) + .projectName(projectName) + .spanCount(i * 3) // e.g., 0, 3, 6, 9, 12 + .usage(null) + .feedbackScores(null) + .endTime(Instant.now()) + .comments(null) + .build()) + .collect(Collectors.toList()); + + int expectedTotalSpanCount = traces.stream().mapToInt(Trace::spanCount).sum(); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + // For each trace, create the actual number of spans matching the spanCount + List allSpans = new ArrayList<>(); + for (Trace trace : traces) { + List spansForTrace = IntStream.range(0, trace.spanCount()) + .mapToObj(j -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(projectName) + .type(SpanType.llm) + .traceId(trace.id()) + .build()) + .toList(); + allSpans.addAll(spansForTrace); + } + spanResourceClient.batchCreateSpans(allSpans, apiKey, workspaceName); + + // Retrieve traces from the API + UUID projectId = getProjectId(projectName, workspaceName, apiKey); + Trace.TracePage resultPage = traceResourceClient.getTraces(projectName, projectId, apiKey, workspaceName, + List.of(), List.of(), 100, Map.of()); + List returnedTraces = resultPage.content(); + + // Check that all created traces are present and have the correct spanCount + for (Trace created : traces) { + returnedTraces.stream() + .filter(returned -> returned.id().equals(created.id())) + .findFirst() + .ifPresentOrElse(returned -> { + assertThat(returned.spanCount()) + .as("Trace with id %s should have spanCount %d", created.id(), created.spanCount()) + .isEqualTo(created.spanCount()); + assertThat(returned.llmSpanCount()) + .as("Trace with id %s should have llmSpanCount %d", created.id(), + created.spanCount()) + .isEqualTo(created.spanCount()); + }, + () -> assertThat(false) + .as("Trace with id %s should be present", created.id()) + .isTrue()); + } + + int actualTotalSpanCount = returnedTraces.stream() + .filter(rt -> traces.stream().anyMatch(t -> t.id().equals(rt.id()))) + .mapToInt(Trace::spanCount) + .sum(); + + assertThat(actualTotalSpanCount) + .as("Total spanCount across all traces should match the expected total") + .isEqualTo(expectedTotalSpanCount); + + int actualTotalLlmSpanCount = returnedTraces.stream() + .filter(rt -> traces.stream().anyMatch(t -> t.id().equals(rt.id()))) + .mapToInt(Trace::llmSpanCount) + .sum(); + + assertThat(actualTotalLlmSpanCount) + .as("Total llmSpanCount across all traces should match the expected total") + .isEqualTo(expectedTotalSpanCount); + } + + private Stream getTracesByProject__whenSortingByValidFields__thenReturnTracesSorted() { + + Comparator inputComparator = Comparator.comparing(trace -> trace.input().toString()); + Comparator outputComparator = Comparator.comparing(trace -> trace.output().toString()); + Comparator metadataComparator = Comparator.comparing(trace -> trace.metadata().toString()); + Comparator tagsComparator = Comparator.comparing(trace -> trace.tags().toString()); + Comparator errorInfoComparator = Comparator.comparing(trace -> trace.errorInfo().toString()); + Comparator usageComparator = Comparator.comparing(trace -> trace.usage().get("completion_tokens")); + + return Stream.of( + Arguments.of(Comparator.comparing(Trace::name), + SortingField.builder().field(SortableFields.NAME).direction(Direction.ASC).build()), + Arguments.of(Comparator.comparing(Trace::name).reversed(), + SortingField.builder().field(SortableFields.NAME).direction(Direction.DESC).build()), + Arguments.of(Comparator.comparing(Trace::startTime), + SortingField.builder().field(SortableFields.START_TIME).direction(Direction.ASC).build()), + Arguments.of(Comparator.comparing(Trace::startTime).reversed(), + SortingField.builder().field(SortableFields.START_TIME).direction(Direction.DESC).build()), + Arguments.of(Comparator.comparing(Trace::endTime), + SortingField.builder().field(SortableFields.END_TIME).direction(Direction.ASC).build()), + Arguments.of(Comparator.comparing(Trace::endTime).reversed(), + SortingField.builder().field(SortableFields.END_TIME).direction(Direction.DESC).build()), + Arguments.of( + Comparator.comparing(Trace::duration) + .thenComparing(Comparator.comparing(Trace::id).reversed()), + SortingField.builder().field(SortableFields.DURATION).direction(Direction.ASC).build()), + Arguments.of( + Comparator.comparing(Trace::duration).reversed() + .thenComparing(Comparator.comparing(Trace::id).reversed()), + SortingField.builder().field(SortableFields.DURATION).direction(Direction.DESC).build()), + Arguments.of(inputComparator, + SortingField.builder().field(SortableFields.INPUT).direction(Direction.ASC).build()), + Arguments.of(inputComparator.reversed(), + SortingField.builder().field(SortableFields.INPUT).direction(Direction.DESC).build()), + Arguments.of(outputComparator, + SortingField.builder().field(SortableFields.OUTPUT).direction(Direction.ASC).build()), + Arguments.of(outputComparator.reversed(), + SortingField.builder().field(SortableFields.OUTPUT).direction(Direction.DESC).build()), + Arguments.of(metadataComparator, + SortingField.builder().field(SortableFields.METADATA).direction(Direction.ASC).build()), + Arguments.of(metadataComparator.reversed(), + SortingField.builder().field(SortableFields.METADATA).direction(Direction.DESC).build()), + Arguments.of(tagsComparator, + SortingField.builder().field(SortableFields.TAGS).direction(Direction.ASC).build()), + Arguments.of(tagsComparator.reversed(), + SortingField.builder().field(SortableFields.TAGS).direction(Direction.DESC).build()), + Arguments.of(Comparator.comparing(Trace::id), + SortingField.builder().field(SortableFields.ID).direction(Direction.ASC).build()), + Arguments.of(Comparator.comparing(Trace::id).reversed(), + SortingField.builder().field(SortableFields.ID).direction(Direction.DESC).build()), + Arguments.of(errorInfoComparator, + SortingField.builder().field(SortableFields.ERROR_INFO).direction(Direction.ASC).build()), + Arguments.of(errorInfoComparator.reversed(), + SortingField.builder().field(SortableFields.ERROR_INFO).direction(Direction.DESC).build()), + Arguments.of(Comparator.comparing(Trace::threadId), SortingField.builder() + .field(SortableFields.THREAD_ID).direction(Direction.ASC).build()), + Arguments.of(Comparator.comparing(Trace::threadId).reversed(), SortingField.builder() + .field(SortableFields.THREAD_ID).direction(Direction.DESC).build()), + Arguments.of(Comparator.comparing(Trace::spanCount) + .thenComparing(Comparator.comparing(Trace::id).reversed()), + SortingField.builder().field(SortableFields.SPAN_COUNT).direction(Direction.ASC).build()), + Arguments.of(Comparator.comparing(Trace::spanCount).reversed() + .thenComparing(Comparator.comparing(Trace::id).reversed()), + SortingField.builder().field(SortableFields.SPAN_COUNT).direction(Direction.DESC).build()), + Arguments.of(Comparator.comparing(Trace::llmSpanCount) + .thenComparing(Comparator.comparing(Trace::id).reversed()), + SortingField.builder().field(SortableFields.LLM_SPAN_COUNT).direction(Direction.ASC) + .build()), + Arguments.of(Comparator.comparing(Trace::llmSpanCount).reversed() + .thenComparing(Comparator.comparing(Trace::id).reversed()), + SortingField.builder().field(SortableFields.LLM_SPAN_COUNT).direction(Direction.DESC) + .build()), + Arguments.of(usageComparator, + SortingField.builder().field("usage.completion_tokens").direction(Direction.ASC).build()), + Arguments.of(usageComparator.reversed(), + SortingField.builder().field("usage.completion_tokens").direction(Direction.DESC).build())); + } + + @Test + void getTracesByProject__whenSortingByInvalidField__thenReturn400() { + var field = RandomStringUtils.secure().nextAlphanumeric(10); + var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( + 400, + "Invalid sorting fields '%s'".formatted(field)); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + + var sortingFields = List.of(SortingField.builder().field(field).direction(Direction.ASC).build()); + var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .queryParam("project_name", projectName) + .queryParam("sorting", + URLEncoder.encode(JsonUtils.writeValueAsString(sortingFields), StandardCharsets.UTF_8)) + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .get(); + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + + var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(actualError).isEqualTo(expectedError); + } + + @ParameterizedTest + @EnumSource(Direction.class) + void getTracesByProject__whenSortingByFeedbackScores__thenReturnTracesSorted(Direction direction) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder() + .projectId(null) + .projectName(projectName) + .usage(null) + .feedbackScores(null) + .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) + .comments(null) + .build()) + .map(trace -> trace.toBuilder() + .duration(trace.startTime().until(trace.endTime(), ChronoUnit.MICROS) / 1000.0) + .build()) + .collect(Collectors.toCollection(ArrayList::new)); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + List scoreForTrace = PodamFactoryUtils.manufacturePojoList(factory, + FeedbackScoreBatchItem.class); + + List allScores = new ArrayList<>(); + for (Trace trace : traces) { + for (FeedbackScoreBatchItem item : scoreForTrace) { + + if (traces.getLast().equals(trace) && scoreForTrace.getFirst().equals(item)) { + continue; + } + + allScores.add(item.toBuilder() + .id(trace.id()) + .projectName(trace.projectName()) + .value(factory.manufacturePojo(BigDecimal.class).abs()) + .build()); + } + } + + traceResourceClient.feedbackScores(allScores, apiKey, workspaceName); + + var sortingField = new SortingField( + "feedback_scores.%s".formatted(scoreForTrace.getFirst().name()), + direction); + + Comparator comparing = Comparator.comparing( + (Trace trace) -> trace.feedbackScores() + .stream() + .filter(score -> score.name().equals(scoreForTrace.getFirst().name())) + .findFirst() + .map(FeedbackScore::value) + .orElse(null), + direction == Direction.ASC + ? Comparator.nullsFirst(Comparator.naturalOrder()) + : Comparator.nullsLast(Comparator.reverseOrder())) + .thenComparing(Comparator.comparing(Trace::id).reversed()); + + var expectedTraces = traces.stream() + .map(trace -> trace.toBuilder() + .feedbackScores(allScores + .stream() + .filter(score -> score.id().equals(trace.id())) + .map(scores -> FeedbackScore.builder() + .name(scores.name()) + .value(scores.value()) + .categoryName(scores.categoryName()) + .source(scores.source()) + .reason(scores.reason()) + .build()) + .toList()) + .build()) + .sorted(comparing) + .toList(); + + List sortingFields = List.of(sortingField); + + getAndAssertPage(workspaceName, projectName, null, List.of(), traces, expectedTraces, List.of(), apiKey, + sortingFields, Set.of()); + } + + @ParameterizedTest + @EnumSource(Trace.TraceField.class) + void getTracesByProject__whenExcludeParamIdDefined__thenReturnSpanExcludingFields(Trace.TraceField field) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + .stream() + .map(trace -> trace.toBuilder().projectName(projectName).build()) + .toList(); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + Map expectedComments = traces + .stream() + .map(trace -> Map.entry(trace.id(), + traceResourceClient.generateAndCreateComment(trace.id(), apiKey, workspaceName, 201))) + .collect(toMap(Map.Entry::getKey, Map.Entry::getValue)); + + traces = traces.stream() + .map(span -> span.toBuilder() + .comments(List.of(expectedComments.get(span.id()))) + .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(span.startTime(), + span.endTime())) + .build()) + .toList(); + + List spans = traces.stream() + .map(trace -> factory.manufacturePojo(Span.class).toBuilder() + .projectName(trace.projectName()) + .traceId(trace.id()) + .build()) + .toList(); + + batchCreateSpansAndAssert(spans, apiKey, workspaceName); + + traces = traces.stream() + .map(trace -> trace.toBuilder() + .totalEstimatedCost(spans.stream() + .filter(span -> span.traceId().equals(trace.id())) + .map(Span::totalEstimatedCost) + .reduce(BigDecimal.ZERO, BigDecimal::add)) + .spanCount((int) spans.stream() + .filter(span -> span.traceId().equals(trace.id())) + .count()) + .usage(spans.stream() + .filter(span -> span.traceId().equals(trace.id())) + .map(Span::usage) + .flatMap(map -> map.entrySet().stream()) + .collect(Collectors.groupingBy(Map.Entry::getKey, + Collectors.summingLong(Map.Entry::getValue)))) + .build()) + .toList(); + + List finalTraces = traces; + List scoreForSpan = IntStream.range(0, traces.size()) + .mapToObj(i -> initFeedbackScoreItem() + .projectName(finalTraces.get(i).projectName()) + .id(finalTraces.get(i).id()) + .build()) + .collect(Collectors.toList()); + + traceResourceClient.feedbackScores(scoreForSpan, apiKey, workspaceName); + + traces = traces.stream() + .map(trace -> trace.toBuilder() + .feedbackScores( + scoreForSpan + .stream() + .filter(score -> score.id().equals(trace.id())) + .map(scores -> FeedbackScore.builder() + .name(scores.name()) + .value(scores.value()) + .categoryName(scores.categoryName()) + .source(scores.source()) + .reason(scores.reason()) + .build()) + .toList()) + .build()) + .toList(); + + List guardrailsByTraceId = traces.stream() + .map(trace -> guardrailsGenerator.generateGuardrailsForTrace(trace.id(), randomUUID(), + trace.projectName())) + .flatMap(Collection::stream) + .toList(); + + traces = traces.stream() + .map(trace -> trace.toBuilder() + .guardrailsValidations( + GuardrailsMapper.INSTANCE.mapToValidations( + guardrailsByTraceId + .stream() + .filter(gr -> gr.entityId().equals(trace.id())) + .toList())) + .build()) + .toList(); + + guardrailsResourceClient.addBatch(guardrailsByTraceId, apiKey, workspaceName); + + traces = traces.stream() + .map(span -> TraceAssertions.EXCLUDE_FUNCTIONS.get(field).apply(span)) + .toList(); + + Set exclude = Set.of(field); + + getAndAssertPage(workspaceName, projectName, null, List.of(), traces, traces.reversed(), List.of(), apiKey, + List.of(), exclude); + + } + + @Test + @DisplayName("should handle filter with percent characters in value correctly") + void shouldHandleTracesWithPercentCharactersInName() { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var traceName = "test%"; + + // Create a trace with % characters in the name + var traces = List.of(createTrace().toBuilder() + .projectName(projectName) + .name(traceName) + .usage(null) + .feedbackScores(null) + .threadId(null) + .comments(null) + .totalEstimatedCost(null) + .build()); + + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + + // Create a filter to search for the trace by name + var filter = TraceFilter.builder() + .field(TraceField.NAME) + .operator(Operator.EQUAL) + .value(traceName) + .build(); + + getAndAssertPage(workspaceName, projectName, null, List.of(filter), traces, traces, List.of(), + apiKey, null, Set.of()); + } + } + + private Integer randomNumber() { + return randomNumber(10, 99); + } + + private static int randomNumber(int minValue, int maxValue) { + return PodamUtils.getIntegerInRange(minValue, maxValue); + } + + private void getAndAssertPage(String workspaceName, String projectName, UUID projectId, + List filters, + List traces, + List expectedTraces, + List unexpectedTraces, + String apiKey, + List sortingFields, + Set exclude) { + int page = 1; + + int size = traces.size() + expectedTraces.size() + unexpectedTraces.size(); + getAndAssertPage(page, size, projectName, projectId, filters, expectedTraces, unexpectedTraces, + workspaceName, apiKey, sortingFields, expectedTraces.size(), exclude); + } + + private void getAndAssertPage(int page, int size, String projectName, UUID projectId, + List filters, + List expectedTraces, List unexpectedTraces, String workspaceName, String apiKey, + List sortingFields, int total, Set exclude) { + + WebTarget target = client.target(URL_TEMPLATE.formatted(baseURI)); + + if (CollectionUtils.isNotEmpty(sortingFields)) { + target = target.queryParam("sorting", + URLEncoder.encode(JsonUtils.writeValueAsString(sortingFields), StandardCharsets.UTF_8)); + } + + if (page > 0) { + target = target.queryParam("page", page); + } + + if (size > 0) { + target = target.queryParam("size", size); + } + + if (projectName != null) { + target = target.queryParam("project_name", projectName); + } + + if (projectId != null) { + target = target.queryParam("project_id", projectId); + } + + if (CollectionUtils.isNotEmpty(exclude)) { + target = target.queryParam("exclude", toURLEncodedQueryParam(List.copyOf(exclude))); + } + + var actualResponse = target + .queryParam("filters", toURLEncodedQueryParam(filters)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get(); + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK); + + var actualPage = actualResponse.readEntity(Trace.TracePage.class); + var actualTraces = actualPage.content(); + + assertThat(actualPage.page()).isEqualTo(page); + assertThat(actualPage.size()).isEqualTo(expectedTraces.size()); + assertThat(actualPage.total()).isEqualTo(total); + + TraceAssertions.assertTraces(actualTraces, expectedTraces, unexpectedTraces, USER); + } + + private List updateFeedbackScore(List feedbackScores, int index, double val) { + feedbackScores.set(index, feedbackScores.get(index).toBuilder() + .value(BigDecimal.valueOf(val)) + .build()); + return feedbackScores; + } + + private List updateFeedbackScore( + List destination, List source, int index) { + destination.set(index, source.get(index).toBuilder().build()); + return destination; + } + + private Trace createTrace() { + return fromBuilder(factory.manufacturePojo(Trace.class).toBuilder()); + } + + private Trace fromBuilder(Trace.TraceBuilder builder) { + return builder + .feedbackScores(null) + .threadId(null) + .comments(null) + .totalEstimatedCost(null) + .usage(null) + .errorInfo(null) + .build(); + } + + private UUID create(Trace trace, String apiKey, String workspaceName) { + return traceResourceClient.createTrace(trace, apiKey, workspaceName); + } + + private void create(UUID entityId, FeedbackScore score, String workspaceName, String apiKey) { + traceResourceClient.feedbackScore(entityId, score, workspaceName, apiKey); + } + + private void batchCreateSpansAndAssert(List expectedSpans, String apiKey, String workspaceName) { + spanResourceClient.batchCreateSpans(expectedSpans, apiKey, workspaceName); + } + + private FeedbackScoreBatchItemBuilder initFeedbackScoreItem() { + return factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder(); + } + + private Map aggregateSpansUsage(List spans) { + if (CollectionUtils.isEmpty(spans)) { + return null; + } + return spans.stream() + .filter(span -> span.usage() != null) + .flatMap(span -> span.usage().entrySet().stream()) + .map(entry -> Map.entry(entry.getKey(), Long.valueOf(entry.getValue()))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, Long::sum)); + } + + private Trace updateSpanCounts(Trace trace, List spans) { + return updateSpanCounts(List.of(trace), spans).getFirst(); + } + + private List updateSpanCounts(List traces, List spans) { + var spansByTraceId = spans.stream().collect(Collectors.groupingBy(Span::traceId)); + return updateSpanCounts(traces, spansByTraceId); + } + + private List updateSpanCounts(List traces, Map> spansByTraceId) { + return traces.stream() + .map(trace -> { + List ts = spansByTraceId.getOrDefault(trace.id(), List.of()); + var total = ts.size(); + var llmCount = ts.stream().filter(s -> s.type() == SpanType.llm).toList().size(); + return trace.toBuilder() + .spanCount(total) + .llmSpanCount(llmCount) + .build(); + }) + .toList(); + } + + @Nested + @DisplayName("Get Traces With Time Filtering:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class GetTracesByProjectTimeFilteringTests { + + private final TraceStatsAssertion traceStatsAssertion = new TraceStatsAssertion(traceResourceClient); + private final TraceTestAssertion traceTestAssertion = new TraceTestAssertion(traceResourceClient, USER); + + // Scenario 1: Boundary condition testing - traces at exact lower bound, upper bound, and in between + private Stream provideBoundaryScenarios() { + return Stream.of( + Arguments.of("/traces/stats", traceStatsAssertion), + Arguments.of("/traces", traceTestAssertion)); + } + + @ParameterizedTest + @DisplayName("filter traces by UUID creation time - includes traces at lower bound, upper bound, and between") + @MethodSource("provideBoundaryScenarios") + void whenTimeParametersProvided_thenIncludeTracesWithinBounds( + String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + Instant baseTime = Instant.now(); + Instant lowerBound = baseTime.minus(Duration.ofMinutes(10)); + Instant upperBound = baseTime; + + // Create traces with UUIDs at specific boundary times + List allTraces = new ArrayList<>(); + + // Index 0: At exact lower bound (should be included) + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(lowerBound)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Index 1: At exact upper bound (should be included) + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(upperBound)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Index 2: Between bounds (should be included) + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(lowerBound.plus(Duration.ofMinutes(5)))) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + traceResourceClient.batchCreateTraces(allTraces, apiKey, workspaceName); + + var queryParams = Map.of( + "from_time", lowerBound.toString(), + "to_time", upperBound.toString()); + + // Clear projectName from traces since API returns projectName=null + allTraces = allTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + + // Sort by ID descending to match API response order + var expectedTraces = allTraces.stream() + .sorted(Comparator.comparing(Trace::id).reversed()) + .toList(); + + var values = testAssertion.transformTestParams(allTraces, expectedTraces, List.of()); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), + values.unexpected(), values.all(), List.of(), queryParams); + } + + @ParameterizedTest + @DisplayName("filter traces by UUID creation time - excludes traces outside bounds") + @MethodSource("provideBoundaryScenarios") + void whenTimeParametersProvided_thenExcludeTracesOutsideBounds( + String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + Instant baseTime = Instant.now(); + Instant lowerBound = baseTime.minus(Duration.ofMinutes(10)); + Instant upperBound = baseTime; + + // Create traces: 3 within bounds, 2 outside bounds + List allTraces = new ArrayList<>(); + + // Within bounds + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(lowerBound)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(upperBound)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(lowerBound.plus(Duration.ofMinutes(1)))) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Outside bounds (before lower) + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(lowerBound.minus(Duration.ofMinutes(1)))) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Outside bounds (after upper) + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(upperBound.plus(Duration.ofMinutes(1)))) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + traceResourceClient.batchCreateTraces(allTraces, apiKey, workspaceName); + + // Expected: indices 0, 1, 2 (within bounds) + // Unexpected: indices 3, 4 (outside bounds) + List expectedTraces = allTraces.subList(0, 3); + List unexpectedTraces = allTraces.subList(3, 5); + + var queryParams = Map.of( + "from_time", lowerBound.toString(), + "to_time", upperBound.toString()); + + // Clear projectName from traces since API returns projectName=null + allTraces = allTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + expectedTraces = expectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + unexpectedTraces = unexpectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + + // Sort expectedTraces by ID descending to match API response order + expectedTraces = expectedTraces.stream() + .sorted(Comparator.comparing(Trace::id).reversed()) + .toList(); + + var values = testAssertion.transformTestParams(allTraces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), + values.unexpected(), values.all(), List.of(), queryParams); + } + + // Scenario 2: ISO-8601 format parsing with extended time range + private Stream provideFormatParsingScenarios() { + return Stream.of( + Arguments.of("/traces/stats", traceStatsAssertion), + Arguments.of("/traces", traceTestAssertion)); + } + + @ParameterizedTest + @DisplayName("time parameters in ISO-8601 format parse correctly and filter traces") + @MethodSource("provideFormatParsingScenarios") + void whenTimeParametersInISO8601Format_thenReturnFilteredTraces( + String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + Instant referenceTime = Instant.now(); + Instant startTime = referenceTime.minus(Duration.ofMinutes(90)); + Instant withinBoundsTime = referenceTime.minus(Duration.ofMinutes(30)); + Instant outsideBoundsTime = referenceTime.plus(Duration.ofMinutes(30)); + + List allTraces = new ArrayList<>(); + + // Should be included: at start of range + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(startTime)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Should be included: within range + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(withinBoundsTime)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Should be included: at end of range + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(referenceTime)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + // Should NOT be included: outside range + allTraces.add(createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(outsideBoundsTime)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build()); + + traceResourceClient.batchCreateTraces(allTraces, apiKey, workspaceName); + + // Filter to get first 3 traces (within bounds) + List expectedTraces = allTraces.subList(0, 3); + List unexpectedTraces = allTraces.subList(3, 4); + + var queryParams = Map.of( + "from_time", startTime.toString(), + "to_time", referenceTime.toString()); + + // Clear projectName from traces since API returns projectName=null + allTraces = allTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + expectedTraces = expectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + unexpectedTraces = unexpectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + + // Sort expectedTraces by ID descending to match API response order + expectedTraces = expectedTraces.stream() + .sorted(Comparator.comparing(Trace::id).reversed()) + .toList(); + + var values = testAssertion.transformTestParams(allTraces, expectedTraces, unexpectedTraces); + + testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), + values.unexpected(), values.all(), List.of(), queryParams); + } + + // Scenario 3: Incomplete time parameters should be rejected + private Stream provideInvalidParameterScenarios() { + return Stream.of( + Arguments.of("/traces/stats", traceStatsAssertion), + Arguments.of("/traces", traceTestAssertion)); + } + + @ParameterizedTest + @DisplayName("time filtering requires both from_time and to_time parameters") + @MethodSource("provideInvalidParameterScenarios") + void whenOnlyFromTimeProvided_thenReturnBadRequest( + String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + Instant now = Instant.now(); + + WebTarget target = client.target(URL_TEMPLATE.formatted(baseURI)) + .queryParam("project_name", projectName) + .queryParam("from_time", now.toString()); + + var actualResponse = target + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get(); + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + } + + @ParameterizedTest + @DisplayName("from_time must be before to_time") + @MethodSource("provideInvalidParameterScenarios") + void whenFromTimeAfterToTime_thenReturnBadRequest( + String endpoint, TracePageTestAssertion testAssertion) { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + Instant now = Instant.now(); + Instant earlier = now.minus(Duration.ofMinutes(10)); + + // from_time (now) is after to_time (earlier) - should fail + WebTarget target = client.target(URL_TEMPLATE.formatted(baseURI)) + .queryParam("project_name", projectName) + .queryParam("from_time", now.toString()) + .queryParam("to_time", earlier.toString()); + + var actualResponse = target + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get(); + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + } + + private UUID generateUUIDForTimestamp(Instant timestamp) { + byte[] zeroBytes = new byte[8]; + return OpenTelemetryMapper.convertOtelIdToUUIDv7(zeroBytes, timestamp.toEpochMilli()); + } + } + +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java index b3019d9ae4c..479a89b7d63 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java @@ -12,7 +12,6 @@ import com.comet.opik.api.FeedbackScoreItem; import com.comet.opik.api.FeedbackScoreItem.FeedbackScoreBatchItem.FeedbackScoreBatchItemBuilder; import com.comet.opik.api.FeedbackScoreNames; -import com.comet.opik.api.Guardrail; import com.comet.opik.api.Project; import com.comet.opik.api.ReactServiceErrorResponse; import com.comet.opik.api.ScoreSource; @@ -30,11 +29,8 @@ import com.comet.opik.api.attachment.EntityType; import com.comet.opik.api.error.ErrorMessage; import com.comet.opik.api.filter.Field; -import com.comet.opik.api.filter.FieldType; import com.comet.opik.api.filter.Filter; import com.comet.opik.api.filter.Operator; -import com.comet.opik.api.filter.TraceField; -import com.comet.opik.api.filter.TraceFilter; import com.comet.opik.api.filter.TraceThreadField; import com.comet.opik.api.filter.TraceThreadFilter; import com.comet.opik.api.resources.utils.AuthTestUtils; @@ -50,8 +46,6 @@ import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.api.resources.utils.resources.AnnotationQueuesResourceClient; import com.comet.opik.api.resources.utils.resources.AttachmentResourceClient; -import com.comet.opik.api.resources.utils.resources.GuardrailsGenerator; -import com.comet.opik.api.resources.utils.resources.GuardrailsResourceClient; import com.comet.opik.api.resources.utils.resources.ProjectResourceClient; import com.comet.opik.api.resources.utils.resources.SpanResourceClient; import com.comet.opik.api.resources.utils.resources.ThreadCommentResourceClient; @@ -59,16 +53,10 @@ import com.comet.opik.api.resources.utils.spans.SpanAssertions; import com.comet.opik.api.resources.utils.traces.TraceAssertions; import com.comet.opik.api.resources.utils.traces.TraceDBUtils; -import com.comet.opik.api.resources.utils.traces.TracePageTestAssertion; -import com.comet.opik.api.resources.utils.traces.TraceStatsAssertion; -import com.comet.opik.api.resources.utils.traces.TraceStreamTestAssertion; -import com.comet.opik.api.resources.utils.traces.TraceTestAssertion; import com.comet.opik.api.sorting.Direction; import com.comet.opik.api.sorting.SortableFields; import com.comet.opik.api.sorting.SortingField; import com.comet.opik.domain.FeedbackScoreMapper; -import com.comet.opik.domain.GuardrailResult; -import com.comet.opik.domain.GuardrailsMapper; import com.comet.opik.domain.SpanType; import com.comet.opik.domain.cost.CostService; import com.comet.opik.domain.filter.FilterQueryBuilder; @@ -135,9 +123,7 @@ import java.util.AbstractMap; import java.util.ArrayList; import java.util.Base64; -import java.util.Collection; import java.util.Comparator; -import java.util.EnumMap; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedHashSet; @@ -147,10 +133,8 @@ import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Function; -import java.util.function.Predicate; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -160,7 +144,6 @@ import static com.comet.opik.api.FeedbackScoreItem.FeedbackScoreBatchItemThread; import static com.comet.opik.api.Visibility.PRIVATE; import static com.comet.opik.api.Visibility.PUBLIC; -import static com.comet.opik.api.filter.TraceField.CUSTOM; import static com.comet.opik.api.resources.utils.ClickHouseContainerUtils.DATABASE_NAME; import static com.comet.opik.api.resources.utils.CommentAssertionUtils.assertComment; import static com.comet.opik.api.resources.utils.CommentAssertionUtils.assertComments; @@ -173,7 +156,6 @@ import static com.comet.opik.api.resources.utils.TestHttpClientUtils.PROJECT_NOT_FOUND_MESSAGE; import static com.comet.opik.api.resources.utils.TestHttpClientUtils.UNAUTHORIZED_RESPONSE; import static com.comet.opik.api.resources.utils.TestUtils.toURLEncodedQueryParam; -import static com.comet.opik.api.resources.utils.traces.TraceAssertions.IGNORED_FIELDS_TRACES; import static com.comet.opik.api.validation.InRangeValidator.MAX_ANALYTICS_DB; import static com.comet.opik.api.validation.InRangeValidator.MAX_ANALYTICS_DB_PRECISION_9; import static com.comet.opik.api.validation.InRangeValidator.MIN_ANALYTICS_DB; @@ -188,7 +170,6 @@ import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo; import static java.util.UUID.randomUUID; import static java.util.function.Predicate.not; -import static java.util.stream.Collectors.toCollection; import static java.util.stream.Collectors.toList; import static java.util.stream.Collectors.toMap; import static org.assertj.core.api.Assertions.assertThat; @@ -207,35 +188,6 @@ class TracesResourceTest { private static final String WORKSPACE_ID = UUID.randomUUID().toString(); private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); - public static final Map> EXCLUDE_FUNCTIONS = new EnumMap<>( - Trace.TraceField.class); - - static { - EXCLUDE_FUNCTIONS.put(Trace.TraceField.NAME, it -> it.toBuilder().name(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.START_TIME, it -> it.toBuilder().startTime(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.END_TIME, it -> it.toBuilder().endTime(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.INPUT, it -> it.toBuilder().input(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.OUTPUT, it -> it.toBuilder().output(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.METADATA, it -> it.toBuilder().metadata(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.TAGS, it -> it.toBuilder().tags(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.USAGE, it -> it.toBuilder().usage(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.ERROR_INFO, it -> it.toBuilder().errorInfo(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.CREATED_AT, it -> it.toBuilder().createdAt(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.CREATED_BY, it -> it.toBuilder().createdBy(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.LAST_UPDATED_BY, it -> it.toBuilder().lastUpdatedBy(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.FEEDBACK_SCORES, it -> it.toBuilder().feedbackScores(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.COMMENTS, it -> it.toBuilder().comments(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.GUARDRAILS_VALIDATIONS, - it -> it.toBuilder().guardrailsValidations(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.SPAN_COUNT, it -> it.toBuilder().spanCount(0).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.LLM_SPAN_COUNT, it -> it.toBuilder().llmSpanCount(0).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.TOTAL_ESTIMATED_COST, - it -> it.toBuilder().totalEstimatedCost(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.THREAD_ID, it -> it.toBuilder().threadId(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.DURATION, it -> it.toBuilder().duration(null).build()); - EXCLUDE_FUNCTIONS.put(Trace.TraceField.VISIBILITY_MODE, it -> it.toBuilder().visibilityMode(null).build()); - } - private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); private final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); private final GenericContainer ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer(); @@ -281,8 +233,6 @@ class TracesResourceTest { private ProjectResourceClient projectResourceClient; private TraceResourceClient traceResourceClient; private SpanResourceClient spanResourceClient; - private GuardrailsResourceClient guardrailsResourceClient; - private GuardrailsGenerator guardrailsGenerator; private ThreadCommentResourceClient threadCommentResourceClient; private AttachmentResourceClient attachmentResourceClient; private AnnotationQueuesResourceClient annotationQueuesResourceClient; @@ -300,10 +250,8 @@ void setUpAll(ClientSupport client) { this.projectResourceClient = new ProjectResourceClient(this.client, baseURI, factory); this.traceResourceClient = new TraceResourceClient(this.client, baseURI); this.spanResourceClient = new SpanResourceClient(this.client, baseURI); - this.guardrailsResourceClient = new GuardrailsResourceClient(client, baseURI); this.threadCommentResourceClient = new ThreadCommentResourceClient(client, baseURI); this.annotationQueuesResourceClient = new AnnotationQueuesResourceClient(client, baseURI); - this.guardrailsGenerator = new GuardrailsGenerator(); this.attachmentResourceClient = new AttachmentResourceClient(client); } @@ -1231,208 +1179,85 @@ private void assertExpectedResponseWithoutABody(boolean expected, Response actua } } - @Nested - @DisplayName("Filters Test:") - @TestInstance(TestInstance.Lifecycle.PER_CLASS) - class FilterTest { + private BigDecimal calculateEstimatedCost(List spans) { + return spans.stream() + .map(span -> CostService.calculateCost(span.model(), span.provider(), span.usage(), null)) + .reduce(BigDecimal.ZERO, BigDecimal::add); + } - private final TraceStatsAssertion traceStatsAssertion = new TraceStatsAssertion(traceResourceClient); - private final TraceTestAssertion traceTestAssertion = new TraceTestAssertion(traceResourceClient, USER); - private final TraceStreamTestAssertion traceStreamTestAssertion = new TraceStreamTestAssertion( - traceResourceClient, USER); + private void assertThreadPage(String projectName, UUID projectId, List expectedThreads, + List filters, Map queryParams, String apiKey, String workspaceName) { + assertThreadPage(projectName, projectId, expectedThreads, filters, queryParams, apiKey, workspaceName, + List.of()); + } - private Stream getFilterTestArguments() { - return Stream.of( - Arguments.of( - "/traces/stats", - traceStatsAssertion), - Arguments.of( - "/traces", - traceTestAssertion), - Arguments.of( - "/traces/search", - traceStreamTestAssertion)); - } + private void assertThreadPage(String projectName, UUID projectId, List expectedThreads, + List filters, Map queryParams, String apiKey, String workspaceName, + List sortingFields) { + var actualPage = traceResourceClient.getTraceThreads(projectId, projectName, apiKey, workspaceName, filters, + sortingFields, queryParams); + var actualTraces = actualPage.content(); - private Stream equalAndNotEqualFilters() { - return Stream.of( - Arguments.of( - "/traces/stats", - Operator.EQUAL, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceStatsAssertion), - Arguments.of( - "/traces", - Operator.EQUAL, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceTestAssertion), - Arguments.of( - "/traces/search", - Operator.EQUAL, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceStreamTestAssertion), - Arguments.of( - "/traces/stats", - Operator.NOT_EQUAL, - (Function, List>) traces -> traces.subList(1, traces.size()), - (Function, List>) traces -> List.of(traces.getFirst()), - traceStatsAssertion), - Arguments.of( - "/traces", - Operator.NOT_EQUAL, - (Function, List>) traces -> traces.subList(1, traces.size()), - (Function, List>) traces -> List.of(traces.getFirst()), - traceTestAssertion), - Arguments.of( - "/traces/search", - Operator.NOT_EQUAL, - (Function, List>) traces -> traces.subList(1, traces.size()), - (Function, List>) traces -> List.of(traces.getFirst()), - traceStreamTestAssertion)); - } + assertThat(actualTraces).hasSize(expectedThreads.size()); + assertThat(actualPage.total()).isEqualTo(expectedThreads.size()); - private Stream getUsageKeyArgs() { - return Stream.of( - Arguments.of( - "/traces/stats", - traceStatsAssertion, - "completion_tokens", - TraceField.USAGE_COMPLETION_TOKENS), - Arguments.of( - "/traces/stats", - traceStatsAssertion, - "prompt_tokens", - TraceField.USAGE_PROMPT_TOKENS), - Arguments.of( - "/traces/stats", - traceStatsAssertion, - "total_tokens", - TraceField.USAGE_TOTAL_TOKENS), - Arguments.of( - "/traces", - traceTestAssertion, - "completion_tokens", - TraceField.USAGE_COMPLETION_TOKENS), - Arguments.of( - "/traces", - traceTestAssertion, - "prompt_tokens", - TraceField.USAGE_PROMPT_TOKENS), - Arguments.of( - "/traces", - traceTestAssertion, - "total_tokens", - TraceField.USAGE_TOTAL_TOKENS), - Arguments.of( - "/traces/search", - traceStreamTestAssertion, - "completion_tokens", - TraceField.USAGE_COMPLETION_TOKENS), - Arguments.of( - "/traces/search", - traceStreamTestAssertion, - "prompt_tokens", - TraceField.USAGE_PROMPT_TOKENS), - Arguments.of( - "/traces/search", - traceStreamTestAssertion, - "total_tokens", - TraceField.USAGE_TOTAL_TOKENS)); + TraceAssertions.assertThreads(expectedThreads, actualTraces); + + for (int i = 0; i < expectedThreads.size(); i++) { + var expectedThread = expectedThreads.get(i); + var actualThread = actualTraces.get(i); + + assertThat(actualThread.createdAt()).isBetween(expectedThread.createdAt(), Instant.now()); + assertThat(actualThread.lastUpdatedAt()) + // Some JVMs can resolve higher than microseconds, such as nanoseconds in the Ubuntu AMD64 JVM + .isBetween(expectedThread.lastUpdatedAt().truncatedTo(ChronoUnit.MICROS), Instant.now()); } + } - private Stream getFeedbackScoresArgs() { - return Stream.of( - Arguments.of( - "/traces/stats", - Operator.EQUAL, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceStatsAssertion), - Arguments.of( - "/traces", - Operator.EQUAL, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceTestAssertion), - Arguments.of( - "/traces/search", - Operator.EQUAL, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceStreamTestAssertion), - Arguments.of( - "/traces/stats", - Operator.NOT_EQUAL, - (Function, List>) traces -> traces.subList(2, traces.size()), - (Function, List>) traces -> traces.subList(0, 2), - traceStatsAssertion), - Arguments.of( - "/traces", - Operator.NOT_EQUAL, - (Function, List>) traces -> traces.subList(2, traces.size()), - (Function, List>) traces -> traces.subList(0, 2), - traceTestAssertion), - Arguments.of( - "/traces/search", - Operator.NOT_EQUAL, - (Function, List>) traces -> traces.subList(2, traces.size()), - (Function, List>) traces -> traces.subList(0, 2), - traceStreamTestAssertion)); - } - - private Stream getDurationArgs() { - Stream arguments = Stream.of( - arguments(Operator.EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 1.0), - arguments(Operator.GREATER_THAN, Duration.ofMillis(8L).toNanos() / 1000, 7.0), - arguments(Operator.GREATER_THAN_EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 1.0), - arguments(Operator.GREATER_THAN_EQUAL, Duration.ofMillis(1L).plusNanos(1000).toNanos() / 1000, 1.0), - arguments(Operator.LESS_THAN, Duration.ofMillis(1L).plusNanos(1).toNanos() / 1000, 2.0), - arguments(Operator.LESS_THAN_EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 1.0), - arguments(Operator.LESS_THAN_EQUAL, Duration.ofMillis(1L).toNanos() / 1000, 2.0)); - - return arguments.flatMap(arg -> Stream.of( - arguments("/traces/stats", traceStatsAssertion, arg.get()[0], - arg.get()[1], arg.get()[2]), - arguments("/traces", traceTestAssertion, arg.get()[0], - arg.get()[1], arg.get()[2]), - arguments("/traces/search", traceStreamTestAssertion, - arg.get()[0], - arg.get()[1], arg.get()[2]))); - } - - private Stream getFilterInvalidOperatorForFieldTypeArgs() { - return filterQueryBuilder.getUnSupportedOperators(TraceField.values()) + private String getValidValue(Field field) { + return switch (field.getType()) { + case STRING, LIST, DICTIONARY, CUSTOM, ENUM, STRING_STATE_DB -> + RandomStringUtils.secure().nextAlphanumeric(10); + case NUMBER, DURATION, FEEDBACK_SCORES_NUMBER -> String.valueOf(randomNumber(1, 10)); + case DATE_TIME, DATE_TIME_STATE_DB -> Instant.now().toString(); + case ERROR_CONTAINER -> ""; + }; + } + + private String getKey(Field field) { + return switch (field.getType()) { + case STRING, NUMBER, DURATION, DATE_TIME, LIST, ENUM, ERROR_CONTAINER, STRING_STATE_DB, DATE_TIME_STATE_DB, + DICTIONARY -> + null; + case FEEDBACK_SCORES_NUMBER, CUSTOM -> RandomStringUtils.secure().nextAlphanumeric(10); + }; + } + + private String getInvalidValue(Field field) { + return switch (field.getType()) { + case STRING, DICTIONARY, CUSTOM, LIST, ENUM, ERROR_CONTAINER, STRING_STATE_DB, DATE_TIME_STATE_DB -> " "; + case NUMBER, DURATION, DATE_TIME, FEEDBACK_SCORES_NUMBER -> RandomStringUtils.secure().nextAlphanumeric(10); + }; + } + + @Nested + @DisplayName("Find trace Threads:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class FindTraceThreads { + + private Stream getUnsupportedOperations() { + return filterQueryBuilder.getUnSupportedOperators(TraceThreadField.values()) .entrySet() .stream() .flatMap(filter -> filter.getValue() .stream() .flatMap(operator -> Stream.of( - Arguments.of("/stats", TraceFilter.builder() - .field(filter.getKey()) - .operator(operator) - .key(getKey(filter.getKey())) - .value(getValidValue(filter.getKey())) - .build()), - Arguments.of("/search", TraceFilter.builder() - .field(filter.getKey()) - .operator(operator) - .key(getKey(filter.getKey())) - .value(getValidValue(filter.getKey())) - .build()), - Arguments.of("", TraceFilter.builder() - .field(filter.getKey()) - .operator(operator) - .key(getKey(filter.getKey())) - .value(getValidValue(filter.getKey())) - .build())))); + Arguments.of(true, filter.getKey(), operator, getValidValue(filter.getKey())), + Arguments.of(false, filter.getKey(), operator, getValidValue(filter.getKey()))))); } private Stream getFilterInvalidValueOrKeyForFieldTypeArgs() { - - Stream filters = filterQueryBuilder.getSupportedOperators(TraceField.values()) + return filterQueryBuilder.getSupportedOperators(TraceThreadField.values()) .entrySet() .stream() .flatMap(filter -> filter.getValue() @@ -1440,13 +1265,13 @@ private Stream getFilterInvalidValueOrKeyForFieldTypeArgs() { .flatMap(operator -> switch (filter.getKey().getType()) { case STRING -> Stream.empty(); case DICTIONARY, FEEDBACK_SCORES_NUMBER -> Stream.of( - TraceFilter.builder() + TraceThreadFilter.builder() .field(filter.getKey()) .operator(operator) .key(null) .value(getValidValue(filter.getKey())) .build(), - TraceFilter.builder() + TraceThreadFilter.builder() .field(filter.getKey()) .operator(operator) // if no value is expected, create an invalid filter by an empty key @@ -1455,351 +1280,352 @@ private Stream getFilterInvalidValueOrKeyForFieldTypeArgs() { : getKey(filter.getKey())) .value(getInvalidValue(filter.getKey())) .build()); - case ERROR_CONTAINER -> Stream.of(); - default -> Stream.of(TraceFilter.builder() + default -> Stream.of(TraceThreadFilter.builder() .field(filter.getKey()) .operator(operator) .value(getInvalidValue(filter.getKey())) .build()); - })); - - return filters.flatMap(filter -> Stream.of( - arguments("/stats", filter), - arguments("", filter), - arguments("/search", filter))); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - @DisplayName("when project name and project id are null, then return bad request") - void whenProjectNameAndIdAreNull__thenReturnBadRequest(String endpoint, TracePageTestAssertion testAssertion) { - - Project project = factory.manufacturePojo(Project.class); - var projectId = projectResourceClient.createProject(project, API_KEY, TEST_WORKSPACE); - - testAssertion.assertTest(null, projectId, API_KEY, TEST_WORKSPACE, List.of(), List.of(), List.of(), - List.of(), Map.of()); + })) + .flatMap(operator -> Stream.of( + Arguments.of(true, operator), + Arguments.of(false, operator))); } - private Instant generateStartTime() { - return Instant.now().minusMillis(randomNumber(1, 1000)); + private Stream getValidFilters() { + return Stream.of( + Arguments.of( + (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() + .field(TraceThreadField.ID) + .operator(Operator.EQUAL) + .value(traces.getFirst().threadId()) + .build(), + (Function, List>) traces -> traces, + (Function, List>) traces -> traces.stream() + .map(trace -> trace.toBuilder() + .threadId(UUID.randomUUID().toString()) + .build()) + .toList()), + Arguments.of( + (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() + .field(TraceThreadField.FIRST_MESSAGE) + .operator(Operator.CONTAINS) + .value(traces.stream().min(Comparator.comparing(Trace::startTime)) + .orElseThrow().input().toString().substring(0, 20)) + .build(), + (Function, List>) traces -> traces, + (Function, List>) traces -> traces), + Arguments.of( + (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() + .field(TraceThreadField.LAST_MESSAGE) + .operator(Operator.CONTAINS) + .value(traces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() + .output().toString().substring(0, 20)) + .build(), + (Function, List>) traces -> traces, + (Function, List>) traces -> traces), + Arguments.of( + (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() + .field(TraceThreadField.DURATION) + .operator(Operator.EQUAL) + .key(null) + .value(DurationUtils.getDurationInMillisWithSubMilliPrecision( + traces.stream().min(Comparator.comparing(Trace::startTime)).get() + .startTime(), + traces.stream().max(Comparator.comparing(Trace::endTime)).get().endTime()) + .toString()) + .build(), + (Function, List>) traces -> traces, + (Function, List>) traces -> traces.stream() + .map(trace -> trace.toBuilder() + .endTime(trace.endTime().plusMillis(100)) + .build()) + .toList()), + Arguments.of( + (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() + .field(TraceThreadField.LAST_UPDATED_AT) + .operator(Operator.EQUAL) + .key(null) + .value(traces.stream().max(Comparator.comparing(Trace::lastUpdatedAt)).get() + .lastUpdatedAt().toString()) + .build(), + (Function, List>) traces -> traces, + (Function, List>) traces -> traces), + Arguments.of( + (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() + .field(TraceThreadField.NUMBER_OF_MESSAGES) + .operator(Operator.EQUAL) + .key(null) + .value(String.valueOf(traces.size() * 2)) + .build(), + (Function, List>) traces -> traces, + (Function, List>) traces -> traces.stream() + .map(trace -> trace.toBuilder() + .threadId(UUID.randomUUID().toString()) + .build()) + .toList())) + .flatMap(args -> Stream.of( + Arguments.of(true, args.get()[0], args.get()[1], args.get()[2]), + Arguments.of(false, args.get()[0], args.get()[1], args.get()[2]))); } @ParameterizedTest - @MethodSource("getFilterTestArguments") - void findWithUsage(String endpoint, TracePageTestAssertion testAssertion) { + @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments") + void findWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) { var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() - .projectName(projectName) - .startTime(generateStartTime()) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(BigDecimal.ZERO) - .guardrailsValidations(null) - .build()) - .toList(); - traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - - var traceIdToSpansMap = traces.stream() - .flatMap(trace -> PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(span -> span.toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .totalEstimatedCost(null) - .build())) - .collect(Collectors.groupingBy(Span::traceId)); - batchCreateSpansAndAssert( - traceIdToSpansMap.values().stream().flatMap(List::stream).toList(), API_KEY, TEST_WORKSPACE); - - traces = traces.stream().map(trace -> trace.toBuilder() - .usage(traceIdToSpansMap.get(trace.id()).stream() - .map(Span::usage) - .flatMap(usage -> usage.entrySet().stream()) - .collect(Collectors.groupingBy( - Map.Entry::getKey, Collectors.summingLong(Map.Entry::getValue)))) - .build()).toList(); - - var traceIdToCommentsMap = traces.stream() - .map(trace -> Pair.of(trace.id(), - IntStream.range(0, 5) - .mapToObj(i -> traceResourceClient.generateAndCreateComment(trace.id(), API_KEY, - TEST_WORKSPACE, 201)) - .toList())) - .collect(Collectors.toMap(Pair::getLeft, Pair::getRight)); - - traces = traces.stream().map(trace -> trace.toBuilder() - .usage(traceIdToSpansMap.get(trace.id()).stream() - .map(Span::usage) - .flatMap(usage -> usage.entrySet().stream()) - .collect(Collectors.groupingBy( - Map.Entry::getKey, Collectors.summingLong(Map.Entry::getValue)))) - .comments(traceIdToCommentsMap.get(trace.id())) - .build()).toList(); - - traces = updateSpanCounts(traces, traceIdToSpansMap); - - var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); - - testAssertion.assertTest(projectName, null, API_KEY, TEST_WORKSPACE, values.expected(), values.unexpected(), - values.all(), List.of(), Map.of()); - } + var threadId = UUID.randomUUID().toString(); - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void findWithoutUsage(String endpoint, TracePageTestAssertion testAssertion) { - var apiKey = UUID.randomUUID().toString(); - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + Trace trace = createTrace(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() + var traces = Stream.of(trace) + .map(it -> it.toBuilder() .projectName(projectName) - .startTime(generateStartTime()) .usage(null) - .feedbackScores(null) - .totalEstimatedCost(BigDecimal.ZERO) - .guardrailsValidations(null) + .input(original) + .output(original) + .threadId(threadId) .build()) .toList(); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var spans = traces.stream() - .flatMap(trace -> PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(span -> span.toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .startTime(trace.startTime()) - .usage(null) - .totalEstimatedCost(null) - .build())) + List spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() + .map(span -> span.toBuilder() + .usage(spanResourceClient.getTokenUsage()) + .model(spanResourceClient.randomModel().toString()) + .provider(spanResourceClient.provider()) + .traceId(traces.getFirst().id()) + .projectName(projectName) + .totalEstimatedCost(null) + .build()) .toList(); - batchCreateSpansAndAssert(spans, apiKey, workspaceName); - traces = updateSpanCounts(traces, spans); + batchCreateSpansAndAssert(spans, API_KEY, TEST_WORKSPACE); - var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), List.of(), Map.of()); - } + var projectId = getProjectId(projectName, TEST_WORKSPACE, API_KEY); - @ParameterizedTest - @MethodSource("getFilterTestArguments") - @DisplayName("when project name is not empty, then return traces by project name") - void whenProjectNameIsNotEmpty__thenReturnTracesByProjectName(String endpoint, - TracePageTestAssertion testAssertion) { + var expectedThreads = List.of(TraceThread.builder() + .firstMessage(expected) + .lastMessage(expected) + .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(trace.startTime(), + trace.endTime())) + .projectId(projectId) + .createdBy(USER) + .startTime(trace.startTime()) + .endTime(trace.endTime()) + .numberOfMessages(traces.size() * 2L) + .id(threadId) + .totalEstimatedCost(calculateEstimatedCost(spans)) + .usage(aggregateSpansUsage(spans)) + .createdAt(trace.createdAt()) + .lastUpdatedAt(trace.lastUpdatedAt()) + .status(TraceThreadStatus.ACTIVE) + .build()); - var projectName = UUID.randomUUID().toString(); + Map queryParams = Map.of("page", "1", "size", "5", "truncate", String.valueOf(truncate)); - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); + assertThreadPage(projectName, null, expectedThreads, List.of(), queryParams, API_KEY, + TEST_WORKSPACE); + } - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + @ParameterizedTest + @MethodSource("getUnsupportedOperations") + void whenFilterUnsupportedOperation__thenReturn400(boolean stream, TraceThreadField field, Operator operator, + String value) { + var filter = TraceThreadFilter.builder() + .field(field) + .operator(operator) + .key(getKey(field)) + .value(value) + .build(); - List traces = new ArrayList<>(); + var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( + HttpStatus.SC_BAD_REQUEST, + "Invalid operator '%s' for field '%s' of type '%s'".formatted( + filter.operator().getQueryParamOperator(), + filter.field().getQueryParamField(), + filter.field().getType())); - for (int i = 0; i < 15; i++) { - Trace trace = createTrace() - .toBuilder() - .projectName(projectName) - .endTime(null) - .duration(null) - .output(null) - .tags(null) - .feedbackScores(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var filters = List.of(filter); - traces.add(trace); - } + try (var actualResponse = !stream + ? findThreads(projectName, filters, API_KEY, TEST_WORKSPACE) + : streamThreadSearch(projectName, null, filters, API_KEY, TEST_WORKSPACE)) { - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); - var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(actualError).isEqualTo(expectedError); + } + } - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), List.of(), Map.of()); + private Response findThreads(String projectName, List<@NotNull TraceThreadFilter> filters, String apiKey, + String testWorkspace) { + return traceResourceClient.getTraceThreads(projectName, apiKey, testWorkspace, filters); } @ParameterizedTest - @MethodSource("getFilterTestArguments") - @DisplayName("when project id is not empty, then return traces by project id") - void whenProjectIdIsNotEmpty__thenReturnTracesByProjectId(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var projectName = UUID.randomUUID().toString(); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + @MethodSource("getFilterInvalidValueOrKeyForFieldTypeArgs") + void whenFilterInvalidValueOrKeyForFieldType__thenReturn400(boolean stream, TraceThreadFilter filter) { + var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( + 400, + "Invalid value '%s' or key '%s' for field '%s' of type '%s'".formatted( + filter.value(), + filter.key(), + filter.field().getQueryParamField(), + filter.field().getType())); - Trace trace = createTrace() - .toBuilder() - .projectName(projectName) - .endTime(null) - .duration(null) - .output(null) - .projectId(null) - .tags(null) - .feedbackScores(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var filters = List.of(filter); - create(trace, apiKey, workspaceName); + try (var actualResponse = !stream + ? findThreads(projectName, filters, API_KEY, TEST_WORKSPACE) + : streamThreadSearch(projectName, null, filters, API_KEY, TEST_WORKSPACE)) { - UUID projectId = getProjectId(projectName, workspaceName, apiKey); + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); - var values = testAssertion.transformTestParams(List.of(), List.of(trace), List.of()); + var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(actualError).isEqualTo(expectedError); + } + } - testAssertion.assertTest(null, projectId, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), List.of(), Map.of()); + private Response streamThreadSearch(String projectName, UUID projectId, + List<@NotNull TraceThreadFilter> filters, String apiKey, String testWorkspace) { + return traceResourceClient.callSearchTraceThreadStream(projectName, projectId, apiKey, testWorkspace, + filters); } @ParameterizedTest - @MethodSource("getFilterTestArguments") - @DisplayName("when filtering by workspace name, then return traces filtered") - void whenFilterWorkspaceName__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - - var workspaceName1 = UUID.randomUUID().toString(); - var workspaceName2 = UUID.randomUUID().toString(); - - var projectName1 = UUID.randomUUID().toString(); + @MethodSource("getValidFilters") + void whenFilterThreads__thenReturnThreadsFiltered( + boolean stream, + Function, TraceThreadFilter> getFilter, + Function, List> getExpectedThreads, + Function, List> getUnexpectedThreads) { + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var threadId = UUID.randomUUID().toString(); + var unexpectedThreadId = UUID.randomUUID().toString(); - var workspaceId1 = UUID.randomUUID().toString(); - var workspaceId2 = UUID.randomUUID().toString(); + var traces = IntStream.range(0, 5) + .mapToObj(it -> { + Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); + return createTrace().toBuilder() + .projectName(projectName) + .usage(null) + .threadId(threadId) + .endTime(now.plus(it, ChronoUnit.MILLIS)) + .startTime(now) + .build(); + }) + .collect(Collectors.toList()); - var apiKey1 = UUID.randomUUID().toString(); - var apiKey2 = UUID.randomUUID().toString(); + traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - mockTargetWorkspace(apiKey1, workspaceName1, workspaceId1); - mockTargetWorkspace(apiKey2, workspaceName2, workspaceId2); + List createTraces = traceResourceClient.getByProjectName(projectName, API_KEY, TEST_WORKSPACE); + List expectedTraces = getExpectedThreads.apply(createTraces); - var traces1 = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName1) + var otherTraces = IntStream.range(0, 5) + .mapToObj(it -> createTrace().toBuilder() + .projectName(projectName) .usage(null) - .threadId(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) - .comments(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) + .threadId(unexpectedThreadId) .build()) - .toList(); + .collect(Collectors.toList()); - var traces2 = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName1) - .usage(null) - .threadId(null) - .feedbackScores(null) - .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .toList(); + List unexpectedTraces = getUnexpectedThreads.apply(otherTraces); + + traceResourceClient.batchCreateTraces(unexpectedTraces, API_KEY, TEST_WORKSPACE); - traceResourceClient.batchCreateTraces(traces1, apiKey1, workspaceName1); - traceResourceClient.batchCreateTraces(traces2, apiKey2, workspaceName2); + var projectId = getProjectId(projectName, TEST_WORKSPACE, API_KEY); + + List expectedThreads = getExpectedThreads(expectedTraces, projectId, threadId, List.of(), + TraceThreadStatus.ACTIVE); - var valueTraces1 = testAssertion.transformTestParams(traces1, traces1.reversed(), List.of()); - var valueTraces2 = testAssertion.transformTestParams(traces2, traces2.reversed(), List.of()); + var filter = getFilter.apply(expectedTraces); + + if (!stream) { + assertThreadPage(projectName, null, expectedThreads, List.of(filter), Map.of(), API_KEY, + TEST_WORKSPACE); + } else { + assertTheadStream(projectName, null, API_KEY, TEST_WORKSPACE, expectedThreads, List.of(filter)); + } + } - testAssertion.assertTest(projectName1, null, apiKey1, workspaceName1, valueTraces1.expected(), - valueTraces1.unexpected(), valueTraces1.all(), List.of(), Map.of()); - testAssertion.assertTest(projectName1, null, apiKey2, workspaceName2, valueTraces2.expected(), - valueTraces2.unexpected(), valueTraces2.all(), List.of(), Map.of()); + private Stream getStatusFilterTestArguments() { + return Stream.of( + Arguments.of(true, TraceThreadStatus.ACTIVE, false), + Arguments.of(true, TraceThreadStatus.INACTIVE, true), + Arguments.of(false, TraceThreadStatus.ACTIVE, false), + Arguments.of(false, TraceThreadStatus.INACTIVE, true)); } @ParameterizedTest - @MethodSource("getFilterTestArguments") - @DisplayName("when traces have cost estimation, then return total cost estimation") - void whenTracesHaveCostEstimation__thenReturnTotalCostEstimation(String endpoint, - TracePageTestAssertion testAssertion) { + @MethodSource("getStatusFilterTestArguments") + @DisplayName("When filtering by thread status, should return only threads with matching status") + void whenFilterByStatus__thenReturnThreadsWithMatchingStatus(boolean stream, TraceThreadStatus filterStatus, + boolean shouldCloseThread) { var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var projectName = UUID.randomUUID().toString(); var workspaceId = UUID.randomUUID().toString(); var apiKey = UUID.randomUUID().toString(); mockTargetWorkspace(apiKey, workspaceName, workspaceId); - List traces = new ArrayList<>(); - - for (int i = 0; i < 5; i++) { + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var threadId = UUID.randomUUID().toString(); - Trace trace = createTrace() - .toBuilder() - .projectName(projectName) - .endTime(null) - .duration(null) - .output(null) - .projectId(null) - .tags(null) - .feedbackScores(null) - .usage(null) - .guardrailsValidations(null) - .totalEstimatedCost(BigDecimal.ZERO) - .build(); - - List spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(span -> span.toBuilder() - .usage(spanResourceClient.getTokenUsage()) - .model(spanResourceClient.randomModel().toString()) - .provider(spanResourceClient.provider()) - .traceId(trace.id()) + // Create traces + var traces = IntStream.range(0, 3) + .mapToObj(it -> { + Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); + return createTrace().toBuilder() .projectName(projectName) - .feedbackScores(null) - .totalEstimatedCost(null) - .build()) - .toList(); - - batchCreateSpansAndAssert(spans, apiKey, workspaceName); - - Trace expectedTrace = trace.toBuilder() - .totalEstimatedCost(calculateEstimatedCost(spans)) - .usage(aggregateSpansUsage(spans)) - .build(); + .usage(null) + .threadId(threadId) + .endTime(now.plus(it, ChronoUnit.MILLIS)) + .startTime(now) + .build(); + }) + .collect(Collectors.toList()); - expectedTrace = updateSpanCounts(expectedTrace, spans); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - traces.add(expectedTrace); + // Close the thread if needed to set its status to INACTIVE + if (shouldCloseThread) { + Mono.delay(Duration.ofMillis(500)).block(); + traceResourceClient.closeTraceThread(threadId, null, projectName, apiKey, workspaceName); } - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var projectId = getProjectId(projectName, workspaceName, apiKey); - UUID projectId = getProjectId(projectName, workspaceName, apiKey); + // Create expected threads with the appropriate status + TraceThreadStatus expectedStatus = shouldCloseThread + ? TraceThreadStatus.INACTIVE + : TraceThreadStatus.ACTIVE; + + List expectedThreads = getExpectedThreads(traces, projectId, threadId, List.of(), + expectedStatus); - var values = testAssertion.transformTestParams(traces, traces.reversed(), List.of()); + // Create filter for the specified status + var statusFilter = TraceThreadFilter.builder() + .field(TraceThreadField.STATUS) + .operator(Operator.EQUAL) + .value(filterStatus.getValue()) + .build(); - testAssertion.assertTest(null, projectId, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), List.of(), Map.of()); + if (!stream) { + // When not streaming, assert the thread page with the status filter + assertThreadPage(null, projectId, expectedThreads, List.of(statusFilter), Map.of(), apiKey, + workspaceName); + } else { + // When streaming, assert the threads with the status filter + assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, List.of(statusFilter)); + } } - @ParameterizedTest - @MethodSource("equalAndNotEqualFilters") - void whenFilterIdAndNameEqual__thenReturnTracesFiltered(String endpoint, - Operator operator, - Function, List> getExpectedTraces, - Function, List> getUnexpectedTraces, - TracePageTestAssertion testAssertion) { + @Test + @DisplayName("When filtering by thread tag, should return only threads with matching tags") + void whenFilterByTags__thenReturnThreadsWithMatchingTags() { var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); var workspaceId = UUID.randomUUID().toString(); @@ -1807,193 +1633,125 @@ void whenFilterIdAndNameEqual__thenReturnTracesFiltered(String endpoint, mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var expectedTraces = getExpectedTraces.apply(traces); - var unexpectedTraces = getUnexpectedTraces.apply(traces); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.ID) - .operator(operator) - .value(traces.getFirst().id().toString()) - .build(), - TraceFilter.builder() - .field(TraceField.NAME) - .operator(operator) - .value(traces.getFirst().name()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var threadId = UUID.randomUUID().toString(); - @ParameterizedTest - @MethodSource("equalAndNotEqualFilters") - void whenFilterByThreadEqual__thenReturnTracesFiltered(String endpoint, - Operator operator, - Function, List> getExpectedTraces, - Function, List> getUnexpectedTraces, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); + // Create traces + var traces = IntStream.range(0, 3) + .mapToObj(it -> { + Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); + return createTrace().toBuilder() + .projectName(projectName) + .usage(null) + .threadId(threadId) + .endTime(now.plus(it, ChronoUnit.MILLIS)) + .startTime(now) + .build(); + }) + .collect(Collectors.toList()); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(UUID.randomUUID().toString()) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); + var projectId = getProjectId(projectName, workspaceName, apiKey); - traces.set(traces.size() - 1, traces.getLast().toBuilder() - .threadId(null) - .build()); + // Wait for thread to be created + Mono.delay(Duration.ofMillis(250)).block(); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var createdThread = traceResourceClient.getTraceThread(threadId, projectId, apiKey, workspaceName); - var expectedTraces = getExpectedTraces.apply(traces); - var unexpectedTraces = getUnexpectedTraces.apply(traces); + // Add tags to the thread + var update = factory.manufacturePojo(TraceThreadUpdate.class); + traceResourceClient.updateThread(update, createdThread.threadModelId(), apiKey, workspaceName, 204); - var filters = List.of( - TraceFilter.builder() - .field(TraceField.THREAD_ID) - .operator(operator) - .value(traces.getFirst().threadId()) - .build()); + List expectedThreads = List.of(createdThread.toBuilder().tags(update.tags()).build()); - var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); + // Create filter for the specified status + var statusFilter = TraceThreadFilter.builder() + .field(TraceThreadField.TAGS) + .operator(Operator.CONTAINS) + .value(update.tags().iterator().next()) + .build(); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); + assertThreadPage(null, projectId, expectedThreads, List.of(statusFilter), Map.of(), apiKey, + workspaceName); + assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, List.of(statusFilter)); } - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterNameEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { + @Test + @DisplayName("When filtering by annotation queue id, should return only threads with matching queue ids") + void whenFilterByAnnotationQueueId__thenReturnThreadsWithMatchingTags() { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); var workspaceId = UUID.randomUUID().toString(); var apiKey = UUID.randomUUID().toString(); mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .totalEstimatedCost(null) - .feedbackScores(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + var project = factory.manufacturePojo(Project.class); + var projectId = projectResourceClient.createProject(project, apiKey, workspaceName); + var threadId = UUID.randomUUID().toString(); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); + // Create traces + var traces = IntStream.range(0, 3) + .mapToObj(it -> { + Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); + return createTrace().toBuilder() + .projectName(project.name()) + .usage(null) + .threadId(threadId) + .endTime(now.plus(it, ChronoUnit.MILLIS)) + .startTime(now) + .build(); + }) + .collect(Collectors.toList()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - List filters = List.of(TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.EQUAL) - .value(traces.getFirst().name().toUpperCase()) - .build()); + // Wait for thread to be created + Mono.delay(Duration.ofMillis(250)).block(); - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + var createdThread = traceResourceClient.getTraceThread(threadId, projectId, apiKey, workspaceName); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } + // Create annotation queue for threads + var annotationQueue = factory.manufacturePojo(AnnotationQueue.class) + .toBuilder() + .projectId(projectId) + .scope(AnnotationQueue.AnnotationScope.THREAD) + .build(); - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterNameStartsWith__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); + annotationQueuesResourceClient.createAnnotationQueueBatch( + new LinkedHashSet<>(List.of(annotationQueue)), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + annotationQueuesResourceClient.addItemsToAnnotationQueue( + annotationQueue.id(), Set.of(createdThread.threadModelId()), apiKey, workspaceName, + HttpStatus.SC_NO_CONTENT); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + List expectedThreads = List.of(createdThread); - var filters = List.of(TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.STARTS_WITH) - .value(traces.getFirst().name().substring(0, traces.getFirst().name().length() - 4).toUpperCase()) - .build()); + // Create filter for the specified status + var statusFilter = TraceThreadFilter.builder() + .field(TraceThreadField.ANNOTATION_QUEUE_IDS) + .operator(Operator.CONTAINS) + .value(annotationQueue.id().toString()) + .build(); - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + assertThreadPage(null, projectId, expectedThreads, List.of(statusFilter), Map.of(), apiKey, + workspaceName); + assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, List.of(statusFilter)); + } - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); + private void assertTheadStream(String projectName, UUID projectId, String apiKey, String workspaceName, + List expectedThreads, List filters) { + var actualThreads = traceResourceClient.searchTraceThreadsStream(projectName, projectId, apiKey, + workspaceName, filters); + TraceAssertions.assertThreads(expectedThreads, actualThreads); } @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterNameEndsWith__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - + @EnumSource(Direction.class) + @DisplayName("When sorting threads by feedback score, then threads are returned in correct order") + void sortThreadsByFeedbackScore_withDirection_thenThreadsReturnedInCorrectOrder(Direction direction) { + // Given var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); var workspaceId = UUID.randomUUID().toString(); var apiKey = UUID.randomUUID().toString(); @@ -2001,4487 +1759,401 @@ void whenFilterNameEndsWith__thenReturnTracesFiltered(String endpoint, TracePage mockTargetWorkspace(apiKey, workspaceName, workspaceId); var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + var project = factory.manufacturePojo(Project.class).toBuilder() + .name(projectName) + .build(); - var filters = List.of(TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.ENDS_WITH) - .value(traces.getFirst().name().substring(3).toUpperCase()) - .build()); + UUID projectId = projectResourceClient.createProject(project, apiKey, workspaceName); - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + // Create threads with different feedback scores + var threadId1 = UUID.randomUUID().toString(); + var threadId2 = UUID.randomUUID().toString(); + var threadId3 = UUID.randomUUID().toString(); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } + // Create traces for threads + Trace trace1 = createTrace().toBuilder() + .threadId(threadId1) + .projectId(projectId) + .projectName(projectName) + .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) + .build(); - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterNameContains__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); + Trace trace2 = createTrace().toBuilder() + .threadId(threadId2) + .projectId(projectId) + .projectName(projectName) + .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) + .build(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + Trace trace3 = createTrace().toBuilder() + .threadId(threadId3) + .projectId(projectId) + .projectName(projectName) + .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) + .build(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) + traceResourceClient.batchCreateTraces(List.of(trace1, trace2, trace3), apiKey, workspaceName); + + // Ensure traces are created with a delay + Mono.delay(Duration.ofMillis(500)).block(); + + // Close the threads to set their status to INACTIVE + traceResourceClient.closeTraceThread(threadId1, null, projectName, apiKey, workspaceName); + traceResourceClient.closeTraceThread(threadId2, null, projectName, apiKey, workspaceName); + traceResourceClient.closeTraceThread(threadId3, null, projectName, apiKey, workspaceName); + + // Add feedback scores with different values + String scoreName = RandomStringUtils.secure().nextAlphanumeric(10); + + List scoreItems = Stream.of(threadId1, threadId2, threadId3) + .map(threadId -> factory.manufacturePojo(FeedbackScoreBatchItemThread.class).toBuilder() + .threadId(threadId) .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) + .name(scoreName) .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + .collect(toList()); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); + Instant now = Instant.now(); + traceResourceClient.threadFeedbackScores(scoreItems, apiKey, workspaceName); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); + // Create feedback scores for expected threads + var feedbackScores = scoreItems.stream() + .collect(Collectors.toMap( + FeedbackScoreItem::threadId, + item -> List.of(createExpectedFeedbackScore(item, now)))); - var filters = List.of(TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.CONTAINS) - .value(traces.getFirst().name().substring(2, traces.getFirst().name().length() - 3).toUpperCase()) - .build()); + // Create expected threads in the correct order based on direction + List expectedThreads = Stream.of( + getExpectedThreads(List.of(trace1), projectId, threadId1, List.of(), TraceThreadStatus.INACTIVE, + feedbackScores.get(threadId1)).getFirst(), + getExpectedThreads(List.of(trace2), projectId, threadId2, List.of(), TraceThreadStatus.INACTIVE, + feedbackScores.get(threadId2)).getFirst(), + getExpectedThreads(List.of(trace3), projectId, threadId3, List.of(), TraceThreadStatus.INACTIVE, + feedbackScores.get(threadId3)).getFirst()) + .sorted(Comparator.comparing(thread -> { + var score = feedbackScores.get(thread.id()).stream() + .filter(fs -> fs.name().equals(scoreName)) + .findFirst() + .orElseThrow(); + return direction == Direction.ASC ? score.value() : score.value().negate(); + })) + .toList(); - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); + // When & Then - Sort by feedback scores and verify using assertThreadPage + var sortingFields = List.of( + SortingField.builder() + .field("feedback_scores." + scoreName) + .direction(direction) + .build()); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); + assertThreadPage(projectName, null, expectedThreads, List.of(), Map.of(), apiKey, workspaceName, + sortingFields); } - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterAnnotationQueueIdContains__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var project = factory.manufacturePojo(Project.class); - var projectId = projectResourceClient.createProject(project, apiKey, workspaceName); - - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(project.name()) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .projectName(project.name()) - .build()); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - // Create annotation queue with items - var queue1 = prepareAnnotationQueue(projectId); - var queue2 = prepareAnnotationQueue(projectId); - annotationQueuesResourceClient.createAnnotationQueueBatch( - new LinkedHashSet<>(List.of(queue1, queue2)), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); - - annotationQueuesResourceClient.addItemsToAnnotationQueue( - queue1.id(), Set.of(traces.getFirst().id()), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); - annotationQueuesResourceClient.addItemsToAnnotationQueue( - queue2.id(), Set.of(traces.get(1).id()), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.ANNOTATION_QUEUE_IDS) - .operator(Operator.CONTAINS) - .value(queue1.id().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(project.name(), null, apiKey, workspaceName, values.expected(), - values.unexpected(), - values.all(), - filters, Map.of()); - } - - private AnnotationQueue prepareAnnotationQueue(UUID projectId) { - return factory.manufacturePojo(AnnotationQueue.class) - .toBuilder() - .projectId(projectId) - .scope(AnnotationQueue.AnnotationScope.TRACE) - .build(); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterNameNotContains__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traceName = generator.generate().toString(); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .name(traceName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .name(generator.generate().toString()) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.NOT_CONTAINS) - .value(traceName.toUpperCase()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("equalAndNotEqualFilters") - void whenFilterStartTimeEqual__thenReturnTracesFiltered(String endpoint, - Operator operator, - Function, List> getExpectedTraces, - Function, List> getUnexpectedTraces, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = getExpectedTraces.apply(traces); - var unexpectedTraces = getUnexpectedTraces.apply(traces); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.START_TIME) - .operator(operator) - .value(traces.getFirst().startTime().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterStartTimeGreaterThan__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .startTime(Instant.now().minusSeconds(60 * 5)) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .startTime(Instant.now().plusSeconds(60 * 5)) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.START_TIME) - .operator(Operator.GREATER_THAN) - .value(Instant.now().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterStartTimeGreaterThanEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .startTime(Instant.now().minusSeconds(60 * 5)) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .startTime(Instant.now()) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.START_TIME) - .operator(Operator.GREATER_THAN_EQUAL) - .value(traces.getFirst().startTime().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterStartTimeLessThan__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .startTime(Instant.now().plusSeconds(60 * 5)) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .startTime(Instant.now().minusSeconds(60 * 5)) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.START_TIME) - .operator(Operator.LESS_THAN) - .value(Instant.now().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterStartTimeLessThanEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .startTime(Instant.now().plusSeconds(60 * 5)) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .startTime(Instant.now().minusSeconds(60 * 5)) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.START_TIME) - .operator(Operator.LESS_THAN_EQUAL) - .value(traces.getFirst().startTime().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterEndTimeEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.END_TIME) - .operator(Operator.EQUAL) - .value(traces.getFirst().endTime().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterInputEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .totalEstimatedCost(null) - .feedbackScores(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.INPUT) - .operator(Operator.EQUAL) - .value(traces.getFirst().input().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterOutputEqual__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .totalEstimatedCost(null) - .feedbackScores(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.OUTPUT) - .operator(Operator.EQUAL) - .value(traces.getFirst().output().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterTotalEstimatedCostGreaterThen__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(spanInStream -> spanInStream.toBuilder() - .projectName(projectName) - .traceId(traces.getFirst().id()) - .usage(Map.of("completion_tokens", Math.abs(factory.manufacturePojo(Integer.class)), - "prompt_tokens", Math.abs(factory.manufacturePojo(Integer.class)))) - .model("gpt-3.5-turbo-1106") - .provider("openai") - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toList()); - - batchCreateSpansAndAssert(spans, apiKey, workspaceName); - - var finalTraces = updateSpanCounts(traces, spans); - var unexpectedTraces = finalTraces.subList(1, traces.size()); - var expectedTrace = finalTraces.getFirst().toBuilder() - .usage(aggregateSpansUsage(spans)) - .totalEstimatedCost(calculateEstimatedCost(spans)) - .build(); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.TOTAL_ESTIMATED_COST) - .operator(Operator.GREATER_THAN) - .value("0") - .build()); - - var values = testAssertion.transformTestParams(finalTraces, List.of(expectedTrace), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("equalAndNotEqualFilters") - void whenFilterTotalEstimatedCostEqual_NotEqual__thenReturnTracesFiltered(String endpoint, - Operator operator, - Function, List> getUnexpectedTraces, // Here we swap the expected and unexpected traces - Function, List> getExpectedTraces, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(spanInStream -> spanInStream.toBuilder() - .projectName(projectName) - .traceId(traces.getFirst().id()) - .usage(Map.of("completion_tokens", Math.abs(factory.manufacturePojo(Integer.class)), - "prompt_tokens", Math.abs(factory.manufacturePojo(Integer.class)))) - .model("gpt-3.5-turbo-1106") - .provider("openai") - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toList()); - - var otherSpans = traces.stream().skip(1) - .flatMap(trace -> PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(span -> span.toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .usage(null) - .model(null) - .totalEstimatedCost(null) - .build())) - .toList(); - - var allSpans = Stream.concat(spans.stream(), otherSpans.stream()).toList(); - batchCreateSpansAndAssert(allSpans, apiKey, workspaceName); - - traces.set(0, traces.getFirst().toBuilder() - .usage(aggregateSpansUsage(spans)) - .totalEstimatedCost(calculateEstimatedCost(spans)) - .build()); - - var finalTraces = updateSpanCounts(traces, allSpans); - var expectedTraces = getExpectedTraces.apply(finalTraces); - var unexpectedTraces = getUnexpectedTraces.apply(finalTraces); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.TOTAL_ESTIMATED_COST) - .operator(operator) - .value("0.00") - .build()); - - var values = testAssertion.transformTestParams(finalTraces, expectedTraces.reversed(), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - Stream whenFilterLlmSpanCountOperator__thenReturnTracesFiltered() { - return getFilterTestArguments().flatMap(args -> Stream.of( - Arguments.of(args.get()[0], args.get()[1], Operator.EQUAL), - Arguments.of(args.get()[0], args.get()[1], Operator.NOT_EQUAL), - Arguments.of(args.get()[0], args.get()[1], Operator.GREATER_THAN), - Arguments.of(args.get()[0], args.get()[1], Operator.GREATER_THAN_EQUAL), - Arguments.of(args.get()[0], args.get()[1], Operator.LESS_THAN), - Arguments.of(args.get()[0], args.get()[1], Operator.LESS_THAN_EQUAL))); - } - - @ParameterizedTest - @MethodSource - void whenFilterLlmSpanCountOperator__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - Operator operator) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> { - var llmSpanCount = RandomUtils.secure().randomInt(1, 7); - return trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .spanCount(llmSpanCount + RandomUtils.secure().randomInt(1, 7)) - .llmSpanCount(llmSpanCount) - .build(); - }) - .toList(); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var spans = traces.stream() - .flatMap(trace -> IntStream.range(0, trace.spanCount()) - .mapToObj(i -> factory.manufacturePojo(Span.class).toBuilder() - .usage(null) - .totalEstimatedCost(null) - .projectName(projectName) - .traceId(trace.id()) - .type(i < trace.llmSpanCount() ? SpanType.llm : SpanType.general) - .build())) - .toList(); - - spanResourceClient.batchCreateSpans(spans, apiKey, workspaceName); - - var llmSpanCountToCompareAgainst = traces.getFirst().llmSpanCount(); - - Predicate matchesFilter = makeLlmSpanCountPredicate(operator, llmSpanCountToCompareAgainst); - Comparator traceIdComparator = Comparator.comparing(Trace::id).reversed(); - - var expectedTraces = traces.stream() - .filter(matchesFilter) - .sorted(traceIdComparator) - .collect(Collectors.toList()); - - var unexpectedTraces = traces.stream() - .filter(matchesFilter.negate()) - .sorted(traceIdComparator) - .collect(Collectors.toList()); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.LLM_SPAN_COUNT) - .operator(operator) - .value(Integer.toString(llmSpanCountToCompareAgainst)) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - Predicate makeLlmSpanCountPredicate(Operator operator, int value) { - switch (operator) { - case Operator.EQUAL : - return trace -> trace.llmSpanCount() == value; - case Operator.NOT_EQUAL : - return trace -> trace.llmSpanCount() != value; - case Operator.GREATER_THAN : - return trace -> trace.llmSpanCount() > value; - case Operator.GREATER_THAN_EQUAL : - return trace -> trace.llmSpanCount() >= value; - case Operator.LESS_THAN : - return trace -> trace.llmSpanCount() < value; - case Operator.LESS_THAN_EQUAL : - return trace -> trace.llmSpanCount() <= value; - default : - throw new IllegalArgumentException("Invalid operator for llm span count filtering: " + operator); - } - } - - @ParameterizedTest - @MethodSource("equalAndNotEqualFilters") - void whenFilterMetadataEqualString__thenReturnTracesFiltered(String endpoint, - Operator operator, - Function, List> getExpectedTraces, - Function, List> getUnexpectedTraces, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - traces.forEach(trace -> create(trace, apiKey, workspaceName)); - var expectedTraces = getExpectedTraces.apply(traces); - var unexpectedTraces = getUnexpectedTraces.apply(traces); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(operator) - .key("$.model[0].version") - .value("OPENAI, CHAT-GPT 4.0") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataEqualNumber__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2023,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.EQUAL) - .key("model[0].year") - .value("2023") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataEqualBoolean__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata( - JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":false,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.EQUAL) - .key("model[0].year") - .value("TRUE") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataEqualNull__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .threadId(null) - .totalEstimatedCost(null) - .feedbackScores(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.EQUAL) - .key("model[0].year") - .value("NULL") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataContainsString__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.CONTAINS) - .key("model[0].version") - .value("CHAT-GPT") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataContainsNumber__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .threadId(null) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":\"two thousand twenty " + - "four\",\"version\":\"OpenAI, Chat-GPT 4.0\"}]}")) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2023,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.CONTAINS) - .key("model[0].year") - .value("02") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataContainsBoolean__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata( - JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":false,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.CONTAINS) - .key("model[0].year") - .value("TRU") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataContainsNull__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"Some " + - "version\"}]}")) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .threadId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.CONTAINS) - .key("model[0].year") - .value("NUL") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataGreaterThanNumber__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2020," + - "\"version\":\"OpenAI, Chat-GPT 4.0\"}]}")) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.GREATER_THAN) - .key("model[0].year") - .value("2023") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataGreaterThanString__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"openAI, " + - "Chat-GPT 4.0\"}]}")) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.GREATER_THAN) - .key("model[0].version") - .value("a") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataGreaterThanBoolean__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"openAI, " + - "Chat-GPT 4.0\"}]}")) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.GREATER_THAN) - .key("model[0].year") - .value("a") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataGreaterThanNull__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"openAI, " + - "Chat-GPT 4.0\"}]}")) - .feedbackScores(null) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.GREATER_THAN) - .key("model[0].year") - .value("a") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataLessThanNumber__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2026," + - "\"version\":\"OpenAI, Chat-GPT 4.0\"}]}")) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(0, traces.getFirst().toBuilder() - .metadata(JsonUtils.getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\"}]}")) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.LESS_THAN) - .key("model[0].year") - .value("2025") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataLessThanString__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"openAI, " + - "Chat-GPT 4.0\"}]}")) - .feedbackScores(null) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.LESS_THAN) - .key("model[0].version") - .value("z") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataLessThanBoolean__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":true,\"version\":\"openAI, " + - "Chat-GPT 4.0\"}]}")) - .feedbackScores(null) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.LESS_THAN) - .key("model[0].year") - .value("z") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterMetadataLessThanNull__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .metadata(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":null,\"version\":\"openAI, " + - "Chat-GPT 4.0\"}]}")) - .feedbackScores(null) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.METADATA) - .operator(Operator.LESS_THAN) - .key("model[0].year") - .value("z") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterTagsContains__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.TAGS) - .operator(Operator.CONTAINS) - .value(traces.getFirst().tags().stream() - .toList() - .get(2) - .substring(0, traces.getFirst().name().length() - 4) - .toUpperCase()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getUsageKeyArgs") - void whenFilterUsageEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - String usageKey, - Field field) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var otherUsageValue = randomNumber(1, 8); - var usageValue = randomNumber(); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(Map.of(usageKey, (long) otherUsageValue)) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toList()); - - traces.set(0, traces.getFirst().toBuilder() - .usage(Map.of(usageKey, (long) usageValue)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var traceIdToSpanMap = traces.stream() - .map(trace -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .usage(Map.of(usageKey, otherUsageValue)) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toMap(Span::traceId, Function.identity())); - traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() - .usage(Map.of(usageKey, usageValue)) - .build()); - batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); - - traces = updateSpanCounts(traces, traceIdToSpanMap.values().stream().toList()); - var expectedTraces = List.of(traces.getFirst()); - var unrelatedTraces = List.of(createTrace()); - - traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(field) - .operator(Operator.EQUAL) - .value(traces.getFirst().usage().get(usageKey).toString()) - .build()); - - var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) - .toList(); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getUsageKeyArgs") - void whenFilterUsageGreaterThan__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - String usageKey, - Field field) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(Map.of(usageKey, 123L)) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(1) - .spanCount(1) - .build()) - .collect(Collectors.toList()); - traces.set(0, traces.getFirst().toBuilder() - .usage(Map.of(usageKey, 456L)) - .build()); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var traceIdToSpanMap = traces.stream() - .map(trace -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .usage(Map.of(usageKey, 123)) - .totalEstimatedCost(null) - .type(SpanType.llm) - .build()) - .collect(Collectors.toMap(Span::traceId, Function.identity())); - traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() - .usage(Map.of(usageKey, 456)) - .build()); - batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); - - var expectedTraces = List.of(traces.getFirst()); - var unrelatedTraces = List.of(createTrace()); - - traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(field) - .operator(Operator.GREATER_THAN) - .value("123") - .build()); - - var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) - .toList(); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getUsageKeyArgs") - void whenFilterUsageGreaterThanEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - String usageKey, - Field field) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(Map.of(usageKey, 123L)) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toList()); - traces.set(0, traces.getFirst().toBuilder() - .usage(Map.of(usageKey, 456L)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var traceIdToSpanMap = traces.stream() - .map(trace -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .usage(Map.of(usageKey, 123)) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toMap(Span::traceId, Function.identity())); - traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() - .usage(Map.of(usageKey, 456)) - .build()); - batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); - - traces = updateSpanCounts(traces, traceIdToSpanMap.values().stream().toList()); - var expectedTraces = List.of(traces.getFirst()); - var unrelatedTraces = List.of(createTrace()); - - traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(field) - .operator(Operator.GREATER_THAN_EQUAL) - .value(traces.getFirst().usage().get(usageKey).toString()) - .build()); - - var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) - .toList(); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getUsageKeyArgs") - void whenFilterUsageLessThan__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - String usageKey, - Field field) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(Map.of(usageKey, 456L)) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toList()); - traces.set(0, traces.getFirst().toBuilder() - .usage(Map.of(usageKey, 123L)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var traceIdToSpanMap = traces.stream() - .map(trace -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .usage(Map.of(usageKey, 456)) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toMap(Span::traceId, Function.identity())); - traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() - .usage(Map.of(usageKey, 123)) - .build()); - batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); - - traces = updateSpanCounts(traces, traceIdToSpanMap.values().stream().toList()); - var expectedTraces = List.of(traces.getFirst()); - var unrelatedTraces = List.of(createTrace()); - - traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(field) - .operator(Operator.LESS_THAN) - .value("456") - .build()); - - var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) - .toList(); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getUsageKeyArgs") - void whenFilterUsageLessThanEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - String usageKey, - Field field) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(Map.of(usageKey, 456L)) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(1) - .spanCount(1) - .build()) - .collect(Collectors.toList()); - traces.set(0, traces.getFirst().toBuilder() - .usage(Map.of(usageKey, 123L)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var traceIdToSpanMap = traces.stream() - .map(trace -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(projectName) - .traceId(trace.id()) - .usage(Map.of(usageKey, 456)) - .totalEstimatedCost(null) - .type(SpanType.llm) - .build()) - .collect(Collectors.toMap(Span::traceId, Function.identity())); - traceIdToSpanMap.put(traces.getFirst().id(), traceIdToSpanMap.get(traces.getFirst().id()).toBuilder() - .usage(Map.of(usageKey, 123)) - .build()); - batchCreateSpansAndAssert(traceIdToSpanMap.values().stream().toList(), apiKey, workspaceName); - - var expectedTraces = List.of(traces.getFirst()); - var unrelatedTraces = List.of(createTrace()); - - traceResourceClient.batchCreateTraces(unrelatedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(field) - .operator(Operator.LESS_THAN_EQUAL) - .value(traces.getFirst().usage().get(usageKey).toString()) - .build()); - - var unexpectedTraces = Stream.of(traces.subList(1, traces.size()), unrelatedTraces).flatMap(List::stream) - .toList(); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFeedbackScoresArgs") - void whenFilterFeedbackScoresEqual__thenReturnTracesFiltered(String endpoint, - Operator operator, - Function, List> getExpectedTraces, - Function, List> getUnexpectedTraces, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .totalEstimatedCost(null) - .feedbackScores(trace.feedbackScores().stream() - .map(feedbackScore -> feedbackScore.toBuilder() - .value(factory.manufacturePojo(BigDecimal.class)) - .build()) - .collect(Collectors.toList())) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(1, traces.get(1).toBuilder() - .feedbackScores( - updateFeedbackScore(traces.get(1).feedbackScores(), traces.getFirst().feedbackScores(), 2)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - traces.forEach(trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var expectedTraces = getExpectedTraces.apply(traces); - var unexpectedTraces = getUnexpectedTraces.apply(traces); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(operator) - .key(traces.getFirst().feedbackScores().get(1).name().toUpperCase()) - .value(traces.getFirst().feedbackScores().get(1).value().toString()) - .build(), - TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(operator) - .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) - .value(traces.getFirst().feedbackScores().get(2).value().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource - void getTracesByProject__whenFilterFeedbackScoresIsEmpty__thenReturnTracesFiltered( - Operator operator, - Function, List> getExpectedTraces, - Function, List> getUnexpectedTraces, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .feedbackScores(trace.feedbackScores().stream() - .map(feedbackScore -> feedbackScore.toBuilder() - .value(factory.manufacturePojo(BigDecimal.class)) - .build()) - .collect(Collectors.toList())) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traces.set(traces.size() - 1, traces.getLast().toBuilder().feedbackScores(null).build()); - traces.forEach(trace1 -> create(trace1, apiKey, workspaceName)); - traces.subList(0, traces.size() - 1).forEach(trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - var expectedTraces = getExpectedTraces.apply(traces); - var unexpectedTraces = getUnexpectedTraces.apply(traces); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(operator) - .key(traces.getFirst().feedbackScores().getFirst().name()) - .value("") - .build()); - var values = testAssertion.transformTestParams(traces, expectedTraces.reversed(), unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - private Stream getTracesByProject__whenFilterFeedbackScoresIsEmpty__thenReturnTracesFiltered() { - return Stream.of( - Arguments.of(Operator.IS_NOT_EMPTY, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceTestAssertion), - Arguments.of(Operator.IS_EMPTY, - (Function, List>) traces -> traces.subList(1, traces.size()), - (Function, List>) traces -> List.of(traces.getFirst()), - traceTestAssertion), - Arguments.of(Operator.IS_NOT_EMPTY, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceStatsAssertion), - Arguments.of(Operator.IS_EMPTY, - (Function, List>) traces -> traces.subList(1, traces.size()), - (Function, List>) traces -> List.of(traces.getFirst()), - traceStatsAssertion), - Arguments.of(Operator.IS_NOT_EMPTY, - (Function, List>) traces -> List.of(traces.getFirst()), - (Function, List>) traces -> traces.subList(1, traces.size()), - traceStreamTestAssertion), - Arguments.of(Operator.IS_EMPTY, - (Function, List>) traces -> traces.subList(1, traces.size()), - (Function, List>) traces -> List.of(traces.getFirst()), - traceStreamTestAssertion)); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterFeedbackScoresGreaterThan__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .totalEstimatedCost(null) - .llmSpanCount(0) - .spanCount(0) - .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() - .map(feedbackScore -> feedbackScore.toBuilder() - .value(factory.manufacturePojo(BigDecimal.class)) - .build()) - .collect(Collectors.toList()), 2, 1234.5678)) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 2345.6789)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - traces.forEach(trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) - .build()); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - unexpectedTraces.forEach( - trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.EQUAL) - .value(traces.getFirst().name()) - .build(), - TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(Operator.GREATER_THAN) - .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) - .value("2345.6788") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterFeedbackScoresGreaterThanEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() - .map(feedbackScore -> feedbackScore.toBuilder() - .value(factory.manufacturePojo(BigDecimal.class)) - .build()) - .collect(Collectors.toList()), 2, 1234.5678)) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 2345.6789)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - traces.forEach(trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) - .build()); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - unexpectedTraces.forEach( - trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(Operator.GREATER_THAN_EQUAL) - .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) - .value(traces.getFirst().feedbackScores().get(2).value().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterFeedbackScoresLessThan__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .comments(null) - .totalEstimatedCost(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() - .map(feedbackScore -> feedbackScore.toBuilder() - .value(factory.manufacturePojo(BigDecimal.class)) - .build()) - .collect(Collectors.toList()), 2, 2345.6789)) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 1234.5678)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - traces.forEach(trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) - .build()); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - unexpectedTraces.forEach( - trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(Operator.LESS_THAN) - .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) - .value("2345.6788") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterFeedbackScoresLessThanEqual__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .totalEstimatedCost(null) - .llmSpanCount(0) - .spanCount(0) - .feedbackScores(updateFeedbackScore(trace.feedbackScores().stream() - .map(feedbackScore -> feedbackScore.toBuilder() - .value(factory.manufacturePojo(BigDecimal.class)) - .build()) - .collect(Collectors.toList()), 2, 2345.6789)) - .guardrailsValidations(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .feedbackScores(updateFeedbackScore(traces.getFirst().feedbackScores(), 2, 1234.5678)) - .build());; - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - traces.forEach(trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectName(RandomStringUtils.secure().nextAlphanumeric(20)) - .projectId(null) - .feedbackScores(PodamFactoryUtils.manufacturePojoList(factory, FeedbackScore.class)) - .build()); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - unexpectedTraces.forEach( - trace -> trace.feedbackScores() - .forEach(feedbackScore -> create(trace.id(), feedbackScore, workspaceName, apiKey))); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.FEEDBACK_SCORES) - .operator(Operator.LESS_THAN_EQUAL) - .key(traces.getFirst().feedbackScores().get(2).name().toUpperCase()) - .value(traces.getFirst().feedbackScores().get(2).value().toString()) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getDurationArgs") - void whenFilterByDuration__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion, - Operator operator, - long end, - double duration) { - String workspaceName = UUID.randomUUID().toString(); - String workspaceId = UUID.randomUUID().toString(); - String apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = generator.generate().toString(); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> { - Instant now = Instant.now(); - return trace.toBuilder() - .projectId(null) - .usage(null) - .projectName(projectName) - .feedbackScores(null) - .threadId(null) - .totalEstimatedCost(null) - .startTime(now) - .endTime(Set.of(Operator.LESS_THAN, Operator.LESS_THAN_EQUAL).contains(operator) - ? Instant.now().plusSeconds(2) - : now.plusNanos(1000)) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build(); - }) - .collect(Collectors.toCollection(ArrayList::new)); - - var start = Instant.now().truncatedTo(ChronoUnit.MILLIS); - traces.set(0, traces.getFirst().toBuilder() - .startTime(start) - .endTime(start.plus(end, ChronoUnit.MICROS)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var expectedTraces = List.of(traces.getFirst()); - - var unexpectedTraces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() - .map(span -> span.toBuilder() - .projectId(null) - .build()) - .toList(); - - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of( - TraceFilter.builder() - .field(TraceField.DURATION) - .operator(operator) - .value(String.valueOf(duration)) - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterInvalidOperatorForFieldTypeArgs") - void whenFilterInvalidOperatorForFieldType__thenReturn400(String path, TraceFilter filter) { - - String errorMessage = filter.field().getType() == FieldType.CUSTOM - ? "Invalid key '%s' for custom filter".formatted(filter.key()) - : "Invalid operator '%s' for field '%s' of type '%s'".formatted( - filter.operator().getQueryParamOperator(), - filter.field().getQueryParamField(), - filter.field().getType()); - - var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( - HttpStatus.SC_BAD_REQUEST, errorMessage); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var filters = List.of(filter); - - Response actualResponse; - if (path.equals("/search")) { - actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path(path) - .request() - .header(HttpHeaders.AUTHORIZATION, API_KEY) - .header(WORKSPACE_HEADER, TEST_WORKSPACE) - .post(Entity.json(TraceSearchStreamRequest.builder() - .projectName(projectName) - .filters(filters) - .build())); - - } else { - - actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path(path) - .queryParam("project_name", projectName) - .queryParam("filters", toURLEncodedQueryParam(filters)) - .request() - .header(HttpHeaders.AUTHORIZATION, API_KEY) - .header(WORKSPACE_HEADER, TEST_WORKSPACE) - .get(); - } - - try (actualResponse) { - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); - - var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); - assertThat(actualError).isEqualTo(expectedError); - } - - } - - @ParameterizedTest - @MethodSource("getFilterInvalidValueOrKeyForFieldTypeArgs") - void whenFilterInvalidValueOrKeyForFieldType__thenReturn400(String path, TraceFilter filter) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( - 400, - "Invalid value '%s' or key '%s' for field '%s' of type '%s'".formatted( - filter.value(), - filter.key(), - filter.field().getQueryParamField(), - filter.field().getType())); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var filters = List.of(filter); - - Response actualResponse; - - if (path.equals("/search")) { - - actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path(path) - .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .post(Entity.json(TraceSearchStreamRequest.builder() - .projectName(projectName) - .filters(filters) - .build())); - - } else { - actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path(path) - .queryParam("project_name", projectName) - .queryParam("filters", toURLEncodedQueryParam(filters)) - .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .get(); - } - - try (actualResponse) { - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); - - var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); - assertThat(actualError).isEqualTo(expectedError); - } - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterGuardrails__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .threadId(null) - .totalEstimatedCost(null) - .feedbackScores(null) - .guardrailsValidations(null) - .comments(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var guardrailsByTraceId = traces.stream() - .collect(Collectors.toMap(Trace::id, trace -> guardrailsGenerator.generateGuardrailsForTrace( - trace.id(), randomUUID(), trace.projectName()))); - - // set the first trace with failed guardrails - guardrailsByTraceId.put(traces.getFirst().id(), guardrailsByTraceId.get(traces.getFirst().id()).stream() - .map(guardrail -> guardrail.toBuilder().result(GuardrailResult.FAILED).build()) - .toList()); - - // set the rest of traces with passed guardrails - traces.subList(1, traces.size()).forEach(trace -> guardrailsByTraceId.put(trace.id(), - guardrailsByTraceId.get(trace.id()).stream() - .map(guardrail -> guardrail.toBuilder() - .result(GuardrailResult.PASSED) - .build()) - .toList())); - - guardrailsByTraceId.values() - .forEach(guardrail -> guardrailsResourceClient.addBatch(guardrail, apiKey, - workspaceName)); - - traces = traces.stream().map(trace -> trace.toBuilder() - .guardrailsValidations(GuardrailsMapper.INSTANCE.mapToValidations( - guardrailsByTraceId.get(trace.id()))) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - // assert failed guardrails - var filtersFailed = List.of( - TraceFilter.builder() - .field(TraceField.GUARDRAILS) - .operator(Operator.EQUAL) - .value(GuardrailResult.FAILED.getResult()) - .build()); - - var valuesFailed = testAssertion.transformTestParams(traces, List.of(traces.getFirst()), - traces.subList(1, traces.size())); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, valuesFailed.expected(), - valuesFailed.unexpected(), valuesFailed.all(), filtersFailed, Map.of()); - - // assert passed guardrails - var filtersPassed = List.of( - TraceFilter.builder() - .field(TraceField.GUARDRAILS) - .operator(Operator.EQUAL) - .value(GuardrailResult.PASSED.getResult()) - .build()); - - var valuesPassed = testAssertion.transformTestParams(traces, traces.subList(1, traces.size()).reversed(), - List.of(traces.getFirst())); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, valuesPassed.expected(), - valuesPassed.unexpected(), valuesPassed.all(), filtersPassed, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterErrorIsNotEmpty__thenReturnTracesFiltered(String endpoint, - TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .errorInfo(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .errorInfo(factory.manufacturePojo(ErrorInfo.class)) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.ERROR_INFO) - .operator(Operator.IS_NOT_EMPTY) - .value("") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - - @ParameterizedTest - @MethodSource("getFilterTestArguments") - void whenFilterErrorIsEmpty__thenReturnTracesFiltered(String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .totalEstimatedCost(null) - .threadId(null) - .guardrailsValidations(null) - .llmSpanCount(0) - .spanCount(0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .errorInfo(null) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - var expectedTraces = List.of(traces.getFirst()); - var unexpectedTraces = List.of(createTrace().toBuilder() - .projectId(null) - .build()); - traceResourceClient.batchCreateTraces(unexpectedTraces, apiKey, workspaceName); - - var filters = List.of(TraceFilter.builder() - .field(TraceField.ERROR_INFO) - .operator(Operator.IS_EMPTY) - .value("") - .build()); - - var values = testAssertion.transformTestParams(traces, expectedTraces, unexpectedTraces); - - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), values.unexpected(), - values.all(), - filters, Map.of()); - } - } - - private BigDecimal calculateEstimatedCost(List spans) { - return spans.stream() - .map(span -> CostService.calculateCost(span.model(), span.provider(), span.usage(), null)) - .reduce(BigDecimal.ZERO, BigDecimal::add); - } - - private void assertThreadPage(String projectName, UUID projectId, List expectedThreads, - List filters, Map queryParams, String apiKey, String workspaceName) { - assertThreadPage(projectName, projectId, expectedThreads, filters, queryParams, apiKey, workspaceName, - List.of()); - } - - private void assertThreadPage(String projectName, UUID projectId, List expectedThreads, - List filters, Map queryParams, String apiKey, String workspaceName, - List sortingFields) { - var actualPage = traceResourceClient.getTraceThreads(projectId, projectName, apiKey, workspaceName, filters, - sortingFields, queryParams); - var actualTraces = actualPage.content(); - - assertThat(actualTraces).hasSize(expectedThreads.size()); - assertThat(actualPage.total()).isEqualTo(expectedThreads.size()); - - TraceAssertions.assertThreads(expectedThreads, actualTraces); - - for (int i = 0; i < expectedThreads.size(); i++) { - var expectedThread = expectedThreads.get(i); - var actualThread = actualTraces.get(i); - - assertThat(actualThread.createdAt()).isBetween(expectedThread.createdAt(), Instant.now()); - assertThat(actualThread.lastUpdatedAt()) - // Some JVMs can resolve higher than microseconds, such as nanoseconds in the Ubuntu AMD64 JVM - .isBetween(expectedThread.lastUpdatedAt().truncatedTo(ChronoUnit.MICROS), Instant.now()); - } - } - - private String getValidValue(Field field) { - return switch (field.getType()) { - case STRING, LIST, DICTIONARY, CUSTOM, ENUM, STRING_STATE_DB -> - RandomStringUtils.secure().nextAlphanumeric(10); - case NUMBER, DURATION, FEEDBACK_SCORES_NUMBER -> String.valueOf(randomNumber(1, 10)); - case DATE_TIME, DATE_TIME_STATE_DB -> Instant.now().toString(); - case ERROR_CONTAINER -> ""; - }; - } - - private String getKey(Field field) { - return switch (field.getType()) { - case STRING, NUMBER, DURATION, DATE_TIME, LIST, ENUM, ERROR_CONTAINER, STRING_STATE_DB, DATE_TIME_STATE_DB, - DICTIONARY -> - null; - case FEEDBACK_SCORES_NUMBER, CUSTOM -> RandomStringUtils.secure().nextAlphanumeric(10); - }; - } - - private String getInvalidValue(Field field) { - return switch (field.getType()) { - case STRING, DICTIONARY, CUSTOM, LIST, ENUM, ERROR_CONTAINER, STRING_STATE_DB, DATE_TIME_STATE_DB -> " "; - case NUMBER, DURATION, DATE_TIME, FEEDBACK_SCORES_NUMBER -> RandomStringUtils.secure().nextAlphanumeric(10); - }; - } - - @Nested - @DisplayName("Find trace Threads:") - @TestInstance(TestInstance.Lifecycle.PER_CLASS) - class FindTraceThreads { - - private Stream getUnsupportedOperations() { - return filterQueryBuilder.getUnSupportedOperators(TraceThreadField.values()) - .entrySet() - .stream() - .flatMap(filter -> filter.getValue() - .stream() - .flatMap(operator -> Stream.of( - Arguments.of(true, filter.getKey(), operator, getValidValue(filter.getKey())), - Arguments.of(false, filter.getKey(), operator, getValidValue(filter.getKey()))))); - } - - private Stream getFilterInvalidValueOrKeyForFieldTypeArgs() { - return filterQueryBuilder.getSupportedOperators(TraceThreadField.values()) - .entrySet() - .stream() - .flatMap(filter -> filter.getValue() - .stream() - .flatMap(operator -> switch (filter.getKey().getType()) { - case STRING -> Stream.empty(); - case DICTIONARY, FEEDBACK_SCORES_NUMBER -> Stream.of( - TraceThreadFilter.builder() - .field(filter.getKey()) - .operator(operator) - .key(null) - .value(getValidValue(filter.getKey())) - .build(), - TraceThreadFilter.builder() - .field(filter.getKey()) - .operator(operator) - // if no value is expected, create an invalid filter by an empty key - .key(Operator.NO_VALUE_OPERATORS.contains(operator) - ? "" - : getKey(filter.getKey())) - .value(getInvalidValue(filter.getKey())) - .build()); - default -> Stream.of(TraceThreadFilter.builder() - .field(filter.getKey()) - .operator(operator) - .value(getInvalidValue(filter.getKey())) - .build()); - })) - .flatMap(operator -> Stream.of( - Arguments.of(true, operator), - Arguments.of(false, operator))); - } - - private Stream getValidFilters() { - return Stream.of( - Arguments.of( - (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() - .field(TraceThreadField.ID) - .operator(Operator.EQUAL) - .value(traces.getFirst().threadId()) - .build(), - (Function, List>) traces -> traces, - (Function, List>) traces -> traces.stream() - .map(trace -> trace.toBuilder() - .threadId(UUID.randomUUID().toString()) - .build()) - .toList()), - Arguments.of( - (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() - .field(TraceThreadField.FIRST_MESSAGE) - .operator(Operator.CONTAINS) - .value(traces.stream().min(Comparator.comparing(Trace::startTime)) - .orElseThrow().input().toString().substring(0, 20)) - .build(), - (Function, List>) traces -> traces, - (Function, List>) traces -> traces), - Arguments.of( - (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() - .field(TraceThreadField.LAST_MESSAGE) - .operator(Operator.CONTAINS) - .value(traces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() - .output().toString().substring(0, 20)) - .build(), - (Function, List>) traces -> traces, - (Function, List>) traces -> traces), - Arguments.of( - (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() - .field(TraceThreadField.DURATION) - .operator(Operator.EQUAL) - .key(null) - .value(DurationUtils.getDurationInMillisWithSubMilliPrecision( - traces.stream().min(Comparator.comparing(Trace::startTime)).get() - .startTime(), - traces.stream().max(Comparator.comparing(Trace::endTime)).get().endTime()) - .toString()) - .build(), - (Function, List>) traces -> traces, - (Function, List>) traces -> traces.stream() - .map(trace -> trace.toBuilder() - .endTime(trace.endTime().plusMillis(100)) - .build()) - .toList()), - Arguments.of( - (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() - .field(TraceThreadField.LAST_UPDATED_AT) - .operator(Operator.EQUAL) - .key(null) - .value(traces.stream().max(Comparator.comparing(Trace::lastUpdatedAt)).get() - .lastUpdatedAt().toString()) - .build(), - (Function, List>) traces -> traces, - (Function, List>) traces -> traces), - Arguments.of( - (Function, TraceThreadFilter>) traces -> TraceThreadFilter.builder() - .field(TraceThreadField.NUMBER_OF_MESSAGES) - .operator(Operator.EQUAL) - .key(null) - .value(String.valueOf(traces.size() * 2)) - .build(), - (Function, List>) traces -> traces, - (Function, List>) traces -> traces.stream() - .map(trace -> trace.toBuilder() - .threadId(UUID.randomUUID().toString()) - .build()) - .toList())) - .flatMap(args -> Stream.of( - Arguments.of(true, args.get()[0], args.get()[1], args.get()[2]), - Arguments.of(false, args.get()[0], args.get()[1], args.get()[2]))); - } - - @ParameterizedTest - @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments") - void findWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) { - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var threadId = UUID.randomUUID().toString(); - - Trace trace = createTrace(); - - var traces = Stream.of(trace) - .map(it -> it.toBuilder() - .projectName(projectName) - .usage(null) - .input(original) - .output(original) - .threadId(threadId) - .build()) - .toList(); - - List spans = PodamFactoryUtils.manufacturePojoList(factory, Span.class).stream() - .map(span -> span.toBuilder() - .usage(spanResourceClient.getTokenUsage()) - .model(spanResourceClient.randomModel().toString()) - .provider(spanResourceClient.provider()) - .traceId(traces.getFirst().id()) - .projectName(projectName) - .totalEstimatedCost(null) - .build()) - .toList(); - - batchCreateSpansAndAssert(spans, API_KEY, TEST_WORKSPACE); - - traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - - var projectId = getProjectId(projectName, TEST_WORKSPACE, API_KEY); - - var expectedThreads = List.of(TraceThread.builder() - .firstMessage(expected) - .lastMessage(expected) - .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(trace.startTime(), - trace.endTime())) - .projectId(projectId) - .createdBy(USER) - .startTime(trace.startTime()) - .endTime(trace.endTime()) - .numberOfMessages(traces.size() * 2L) - .id(threadId) - .totalEstimatedCost(calculateEstimatedCost(spans)) - .usage(aggregateSpansUsage(spans)) - .createdAt(trace.createdAt()) - .lastUpdatedAt(trace.lastUpdatedAt()) - .status(TraceThreadStatus.ACTIVE) - .build()); - - Map queryParams = Map.of("page", "1", "size", "5", "truncate", String.valueOf(truncate)); - - assertThreadPage(projectName, null, expectedThreads, List.of(), queryParams, API_KEY, - TEST_WORKSPACE); - } - - @ParameterizedTest - @MethodSource("getUnsupportedOperations") - void whenFilterUnsupportedOperation__thenReturn400(boolean stream, TraceThreadField field, Operator operator, - String value) { - var filter = TraceThreadFilter.builder() - .field(field) - .operator(operator) - .key(getKey(field)) - .value(value) - .build(); - - var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( - HttpStatus.SC_BAD_REQUEST, - "Invalid operator '%s' for field '%s' of type '%s'".formatted( - filter.operator().getQueryParamOperator(), - filter.field().getQueryParamField(), - filter.field().getType())); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var filters = List.of(filter); - - try (var actualResponse = !stream - ? findThreads(projectName, filters, API_KEY, TEST_WORKSPACE) - : streamThreadSearch(projectName, null, filters, API_KEY, TEST_WORKSPACE)) { - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); - - var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); - assertThat(actualError).isEqualTo(expectedError); - } - } - - private Response findThreads(String projectName, List<@NotNull TraceThreadFilter> filters, String apiKey, - String testWorkspace) { - return traceResourceClient.getTraceThreads(projectName, apiKey, testWorkspace, filters); - } - - @ParameterizedTest - @MethodSource("getFilterInvalidValueOrKeyForFieldTypeArgs") - void whenFilterInvalidValueOrKeyForFieldType__thenReturn400(boolean stream, TraceThreadFilter filter) { - var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( - 400, - "Invalid value '%s' or key '%s' for field '%s' of type '%s'".formatted( - filter.value(), - filter.key(), - filter.field().getQueryParamField(), - filter.field().getType())); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var filters = List.of(filter); - - try (var actualResponse = !stream - ? findThreads(projectName, filters, API_KEY, TEST_WORKSPACE) - : streamThreadSearch(projectName, null, filters, API_KEY, TEST_WORKSPACE)) { - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); - - var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); - assertThat(actualError).isEqualTo(expectedError); - } - } - - private Response streamThreadSearch(String projectName, UUID projectId, - List<@NotNull TraceThreadFilter> filters, String apiKey, String testWorkspace) { - return traceResourceClient.callSearchTraceThreadStream(projectName, projectId, apiKey, testWorkspace, - filters); - } - - @ParameterizedTest - @MethodSource("getValidFilters") - void whenFilterThreads__thenReturnThreadsFiltered( - boolean stream, - Function, TraceThreadFilter> getFilter, - Function, List> getExpectedThreads, - Function, List> getUnexpectedThreads) { - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var threadId = UUID.randomUUID().toString(); - var unexpectedThreadId = UUID.randomUUID().toString(); - - var traces = IntStream.range(0, 5) - .mapToObj(it -> { - Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); - return createTrace().toBuilder() - .projectName(projectName) - .usage(null) - .threadId(threadId) - .endTime(now.plus(it, ChronoUnit.MILLIS)) - .startTime(now) - .build(); - }) - .collect(Collectors.toList()); - - traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - - List createTraces = traceResourceClient.getByProjectName(projectName, API_KEY, TEST_WORKSPACE); - List expectedTraces = getExpectedThreads.apply(createTraces); - - var otherTraces = IntStream.range(0, 5) - .mapToObj(it -> createTrace().toBuilder() - .projectName(projectName) - .usage(null) - .threadId(unexpectedThreadId) - .build()) - .collect(Collectors.toList()); - - List unexpectedTraces = getUnexpectedThreads.apply(otherTraces); - - traceResourceClient.batchCreateTraces(unexpectedTraces, API_KEY, TEST_WORKSPACE); - - var projectId = getProjectId(projectName, TEST_WORKSPACE, API_KEY); - - List expectedThreads = getExpectedThreads(expectedTraces, projectId, threadId, List.of(), - TraceThreadStatus.ACTIVE); - - var filter = getFilter.apply(expectedTraces); - - if (!stream) { - assertThreadPage(projectName, null, expectedThreads, List.of(filter), Map.of(), API_KEY, - TEST_WORKSPACE); - } else { - assertTheadStream(projectName, null, API_KEY, TEST_WORKSPACE, expectedThreads, List.of(filter)); - } - } - - private Stream getStatusFilterTestArguments() { - return Stream.of( - Arguments.of(true, TraceThreadStatus.ACTIVE, false), - Arguments.of(true, TraceThreadStatus.INACTIVE, true), - Arguments.of(false, TraceThreadStatus.ACTIVE, false), - Arguments.of(false, TraceThreadStatus.INACTIVE, true)); - } - - @ParameterizedTest - @MethodSource("getStatusFilterTestArguments") - @DisplayName("When filtering by thread status, should return only threads with matching status") - void whenFilterByStatus__thenReturnThreadsWithMatchingStatus(boolean stream, TraceThreadStatus filterStatus, - boolean shouldCloseThread) { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var threadId = UUID.randomUUID().toString(); - - // Create traces - var traces = IntStream.range(0, 3) - .mapToObj(it -> { - Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); - return createTrace().toBuilder() - .projectName(projectName) - .usage(null) - .threadId(threadId) - .endTime(now.plus(it, ChronoUnit.MILLIS)) - .startTime(now) - .build(); - }) - .collect(Collectors.toList()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - // Close the thread if needed to set its status to INACTIVE - if (shouldCloseThread) { - Mono.delay(Duration.ofMillis(500)).block(); - traceResourceClient.closeTraceThread(threadId, null, projectName, apiKey, workspaceName); - } - - var projectId = getProjectId(projectName, workspaceName, apiKey); - - // Create expected threads with the appropriate status - TraceThreadStatus expectedStatus = shouldCloseThread - ? TraceThreadStatus.INACTIVE - : TraceThreadStatus.ACTIVE; - - List expectedThreads = getExpectedThreads(traces, projectId, threadId, List.of(), - expectedStatus); - - // Create filter for the specified status - var statusFilter = TraceThreadFilter.builder() - .field(TraceThreadField.STATUS) - .operator(Operator.EQUAL) - .value(filterStatus.getValue()) - .build(); - - if (!stream) { - // When not streaming, assert the thread page with the status filter - assertThreadPage(null, projectId, expectedThreads, List.of(statusFilter), Map.of(), apiKey, - workspaceName); - } else { - // When streaming, assert the threads with the status filter - assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, List.of(statusFilter)); - } - } - - @Test - @DisplayName("When filtering by thread tag, should return only threads with matching tags") - void whenFilterByTags__thenReturnThreadsWithMatchingTags() { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var threadId = UUID.randomUUID().toString(); - - // Create traces - var traces = IntStream.range(0, 3) - .mapToObj(it -> { - Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); - return createTrace().toBuilder() - .projectName(projectName) - .usage(null) - .threadId(threadId) - .endTime(now.plus(it, ChronoUnit.MILLIS)) - .startTime(now) - .build(); - }) - .collect(Collectors.toList()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var projectId = getProjectId(projectName, workspaceName, apiKey); - - // Wait for thread to be created - Mono.delay(Duration.ofMillis(250)).block(); - - var createdThread = traceResourceClient.getTraceThread(threadId, projectId, apiKey, workspaceName); - - // Add tags to the thread - var update = factory.manufacturePojo(TraceThreadUpdate.class); - traceResourceClient.updateThread(update, createdThread.threadModelId(), apiKey, workspaceName, 204); - - List expectedThreads = List.of(createdThread.toBuilder().tags(update.tags()).build()); - - // Create filter for the specified status - var statusFilter = TraceThreadFilter.builder() - .field(TraceThreadField.TAGS) - .operator(Operator.CONTAINS) - .value(update.tags().iterator().next()) - .build(); - - assertThreadPage(null, projectId, expectedThreads, List.of(statusFilter), Map.of(), apiKey, - workspaceName); - assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, List.of(statusFilter)); - } - - @Test - @DisplayName("When filtering by annotation queue id, should return only threads with matching queue ids") - void whenFilterByAnnotationQueueId__thenReturnThreadsWithMatchingTags() { - - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var project = factory.manufacturePojo(Project.class); - var projectId = projectResourceClient.createProject(project, apiKey, workspaceName); - var threadId = UUID.randomUUID().toString(); - - // Create traces - var traces = IntStream.range(0, 3) - .mapToObj(it -> { - Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); - return createTrace().toBuilder() - .projectName(project.name()) - .usage(null) - .threadId(threadId) - .endTime(now.plus(it, ChronoUnit.MILLIS)) - .startTime(now) - .build(); - }) - .collect(Collectors.toList()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - // Wait for thread to be created - Mono.delay(Duration.ofMillis(250)).block(); - - var createdThread = traceResourceClient.getTraceThread(threadId, projectId, apiKey, workspaceName); - - // Create annotation queue for threads - var annotationQueue = factory.manufacturePojo(AnnotationQueue.class) - .toBuilder() - .projectId(projectId) - .scope(AnnotationQueue.AnnotationScope.THREAD) - .build(); - - annotationQueuesResourceClient.createAnnotationQueueBatch( - new LinkedHashSet<>(List.of(annotationQueue)), apiKey, workspaceName, HttpStatus.SC_NO_CONTENT); - - annotationQueuesResourceClient.addItemsToAnnotationQueue( - annotationQueue.id(), Set.of(createdThread.threadModelId()), apiKey, workspaceName, - HttpStatus.SC_NO_CONTENT); - - List expectedThreads = List.of(createdThread); - - // Create filter for the specified status - var statusFilter = TraceThreadFilter.builder() - .field(TraceThreadField.ANNOTATION_QUEUE_IDS) - .operator(Operator.CONTAINS) - .value(annotationQueue.id().toString()) - .build(); - - assertThreadPage(null, projectId, expectedThreads, List.of(statusFilter), Map.of(), apiKey, - workspaceName); - assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, List.of(statusFilter)); - } - - private void assertTheadStream(String projectName, UUID projectId, String apiKey, String workspaceName, - List expectedThreads, List filters) { - var actualThreads = traceResourceClient.searchTraceThreadsStream(projectName, projectId, apiKey, - workspaceName, filters); - TraceAssertions.assertThreads(expectedThreads, actualThreads); - } - - @ParameterizedTest - @EnumSource(Direction.class) - @DisplayName("When sorting threads by feedback score, then threads are returned in correct order") - void sortThreadsByFeedbackScore_withDirection_thenThreadsReturnedInCorrectOrder(Direction direction) { - // Given - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var project = factory.manufacturePojo(Project.class).toBuilder() - .name(projectName) - .build(); - - UUID projectId = projectResourceClient.createProject(project, apiKey, workspaceName); - - // Create threads with different feedback scores - var threadId1 = UUID.randomUUID().toString(); - var threadId2 = UUID.randomUUID().toString(); - var threadId3 = UUID.randomUUID().toString(); - - // Create traces for threads - Trace trace1 = createTrace().toBuilder() - .threadId(threadId1) - .projectId(projectId) - .projectName(projectName) - .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) - .build(); - - Trace trace2 = createTrace().toBuilder() - .threadId(threadId2) - .projectId(projectId) - .projectName(projectName) - .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) - .build(); - - Trace trace3 = createTrace().toBuilder() - .threadId(threadId3) - .projectId(projectId) - .projectName(projectName) - .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) - .build(); - - traceResourceClient.batchCreateTraces(List.of(trace1, trace2, trace3), apiKey, workspaceName); - - // Ensure traces are created with a delay - Mono.delay(Duration.ofMillis(500)).block(); - - // Close the threads to set their status to INACTIVE - traceResourceClient.closeTraceThread(threadId1, null, projectName, apiKey, workspaceName); - traceResourceClient.closeTraceThread(threadId2, null, projectName, apiKey, workspaceName); - traceResourceClient.closeTraceThread(threadId3, null, projectName, apiKey, workspaceName); - - // Add feedback scores with different values - String scoreName = RandomStringUtils.secure().nextAlphanumeric(10); - - List scoreItems = Stream.of(threadId1, threadId2, threadId3) - .map(threadId -> factory.manufacturePojo(FeedbackScoreBatchItemThread.class).toBuilder() - .threadId(threadId) - .projectName(projectName) - .name(scoreName) - .build()) - .collect(toList()); - - Instant now = Instant.now(); - traceResourceClient.threadFeedbackScores(scoreItems, apiKey, workspaceName); - - // Create feedback scores for expected threads - var feedbackScores = scoreItems.stream() - .collect(Collectors.toMap( - FeedbackScoreItem::threadId, - item -> List.of(createExpectedFeedbackScore(item, now)))); - - // Create expected threads in the correct order based on direction - List expectedThreads = Stream.of( - getExpectedThreads(List.of(trace1), projectId, threadId1, List.of(), TraceThreadStatus.INACTIVE, - feedbackScores.get(threadId1)).getFirst(), - getExpectedThreads(List.of(trace2), projectId, threadId2, List.of(), TraceThreadStatus.INACTIVE, - feedbackScores.get(threadId2)).getFirst(), - getExpectedThreads(List.of(trace3), projectId, threadId3, List.of(), TraceThreadStatus.INACTIVE, - feedbackScores.get(threadId3)).getFirst()) - .sorted(Comparator.comparing(thread -> { - var score = feedbackScores.get(thread.id()).stream() - .filter(fs -> fs.name().equals(scoreName)) - .findFirst() - .orElseThrow(); - return direction == Direction.ASC ? score.value() : score.value().negate(); - })) - .toList(); - - // When & Then - Sort by feedback scores and verify using assertThreadPage - var sortingFields = List.of( - SortingField.builder() - .field("feedback_scores." + scoreName) - .direction(direction) - .build()); - - assertThreadPage(projectName, null, expectedThreads, List.of(), Map.of(), apiKey, workspaceName, - sortingFields); - } - - private Stream getFeedbackScoreFilterTestArguments() { - return Stream.of( - Arguments.of( - true, - Operator.EQUAL, - generateExpectedIndices(), - (BiFunction>) this::generateExpectedEqualsMatch, - (BiFunction>) this::generateUnexpectedEqualsMatch, - (Function) BigDecimal::toString), // Filter value function for EQUAL - Arguments.of( - false, - Operator.EQUAL, - generateExpectedIndices(), - (BiFunction>) this::generateExpectedEqualsMatch, - (BiFunction>) this::generateUnexpectedEqualsMatch, - (Function) BigDecimal::toString), - Arguments.of( - true, - Operator.IS_NOT_EMPTY, - generateExpectedIndices(), - (BiFunction>) (name, - value) -> generateNotEmptyMatch(name), - (BiFunction>) (name, - value) -> generateUnexpectedNotEmptyMatch(name), - (Function) value -> ""), // Empty value for IS_NOT_EMPTY - Arguments.of( - false, - Operator.IS_NOT_EMPTY, - generateExpectedIndices(), - (BiFunction>) (name, - value) -> generateNotEmptyMatch(name), - (BiFunction>) (name, - value) -> generateUnexpectedNotEmptyMatch(name), - (Function) value -> ""), - Arguments.of( - true, - Operator.IS_EMPTY, - generateExpectedIndices(), - (BiFunction>) (name, - value) -> generateIsEmptyMatch(name), - (BiFunction>) (name, - value) -> generateNotEmptyMatch(name), - (Function) value -> ""), // Empty value for IS_EMPTY - Arguments.of( - false, - Operator.IS_EMPTY, - generateExpectedIndices(), - (BiFunction>) (name, - value) -> generateIsEmptyMatch(name), - (BiFunction>) (name, - value) -> generateNotEmptyMatch(name), - (Function) value -> "") - - ); - } - - private Set generateExpectedIndices() { - return new HashSet<>(List.of(RandomUtils.secure().randomInt(0, 5), RandomUtils.secure().randomInt(0, 5))); - } - - private List generateIsEmptyMatch(String name) { - return PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItemThread.class) - .stream() - .filter(score -> !score.name().equals(name)) - .toList(); - } - - private List generateNotEmptyMatch(String name) { - List scores = PodamFactoryUtils.manufacturePojoList(factory, - FeedbackScoreBatchItemThread.class); - - scores.set(0, scores.getFirst().toBuilder() - .name(name) - .build()); - - return scores; - } - - private List generateUnexpectedNotEmptyMatch(String name) { - return PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItemThread.class) - .stream() - .filter(score -> !score.name().equals(name)) - .toList(); - } - - private List generateExpectedEqualsMatch(String name, BigDecimal value) { - List scores = PodamFactoryUtils.manufacturePojoList(factory, - FeedbackScoreBatchItemThread.class); - - scores.set(0, scores.getFirst().toBuilder() - .name(name) - .value(value) - .build()); - - return scores; - } - - private List generateUnexpectedEqualsMatch(String name, BigDecimal value) { - List scores = PodamFactoryUtils.manufacturePojoList(factory, - FeedbackScoreBatchItemThread.class); - - scores.set(0, scores.getFirst().toBuilder() - .name(name) - .build()); - - return scores; - } - - @ParameterizedTest - @MethodSource("getFeedbackScoreFilterTestArguments") - @DisplayName("When filtering by feedback score with different operators, should return matching threads") - void whenFilterByFeedbackScore__thenReturnThreadsWithMatchingFeedbackScore( - boolean stream, - Operator operator, - Set expectedThreadIndices, - BiFunction> matchingScoreFunction, - BiFunction> unmatchingScoreFunction, - Function filterValueFunction) { - - // Given - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var project = factory.manufacturePojo(Project.class).toBuilder() - .name(projectName) - .build(); - - UUID projectId = projectResourceClient.createProject(project, apiKey, workspaceName); - - // Create threads with different feedback scores - var allThreadIds = PodamFactoryUtils.manufacturePojoList(factory, UUID.class); - - // Create traces for threads - var allTraces = allThreadIds - .stream() - .map(threadId -> createTrace().toBuilder() - .threadId(threadId.toString()) - .projectId(null) - .projectName(projectName) - .startTime(Instant.now().minusSeconds(3).truncatedTo(ChronoUnit.MICROS)) - .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) - .build()) - .toList(); - - allTraces.forEach(trace -> traceResourceClient.createTrace(trace, apiKey, workspaceName)); - - // Add feedback scores with different values - Map threadIdAndLastUpdateAts = new HashMap<>(); - - Mono.delay(Duration.ofMillis(500)).block(); - allThreadIds.forEach(threadId -> { - threadIdAndLastUpdateAts.put(threadId.toString(), Instant.now()); - traceResourceClient.closeTraceThread(threadId.toString(), null, projectName, apiKey, workspaceName); - }); - - String targetScoreName = RandomStringUtils.secure().nextAlphanumeric(30); - BigDecimal targetScoreValue = factory.manufacturePojo(BigDecimal.class); - - // Create feedback scores based on the provided map - List expectedScores = allThreadIds - .stream() - .filter(threadId -> isExpected(expectedThreadIndices, threadId, allThreadIds)) - .flatMap(threadId -> { - return matchingScoreFunction.apply(targetScoreName, targetScoreValue).stream() - .map(item -> item.toBuilder() - .threadId(threadId.toString()) - .projectName(projectName) - .build()); - }).collect(Collectors.toList()); - - List unexpectedScores = allThreadIds.stream() - .filter(threadId -> !isExpected(expectedThreadIndices, threadId, allThreadIds)) - .flatMap(threadId -> { - return unmatchingScoreFunction.apply(targetScoreName, targetScoreValue).stream() - .map(item -> item.toBuilder() - .threadId(threadId.toString()) - .projectName(projectName) - .build()); - }).collect(Collectors.toList()); - - List scoreItems = Stream - .concat(expectedScores.stream(), unexpectedScores.stream()) - .toList(); - - // Create feedback scores for threads - Instant feedbackScoreCreationTime = Instant.now(); - traceResourceClient.threadFeedbackScores(scoreItems, apiKey, workspaceName); - - // Determine expected threads based on indices - var expectedThreadIds = allThreadIds.reversed() - .stream() - .filter(threadId -> isExpected(expectedThreadIndices, threadId, allThreadIds)) - .map(UUID::toString) - .toList(); - - // Create expected threads with ALL feedback scores from matching threads - Comparator comparing = Comparator - .comparing((TraceThread traceThread) -> threadIdAndLastUpdateAts.get(traceThread.id())).reversed(); - - List expectedThreads = expectedThreadIds.stream() - .map(threadId -> { - // Get ALL feedback scores for this thread (both expected and unexpected) - var allFeedbackScoresForThread = scoreItems.stream() - .filter(item -> item.threadId().equals(threadId)) - .map(item -> createExpectedFeedbackScore(item, feedbackScoreCreationTime)) - .toList(); - - var traces = allTraces.stream() - .filter(trace -> trace.threadId().equals(threadId)) - .toList(); - - return getExpectedThreads(traces, projectId, threadId, List.of(), - TraceThreadStatus.INACTIVE, allFeedbackScoresForThread).getFirst(); - }) - .sorted(comparing) - .toList(); - - // Create filter for the specific feedback score - var feedbackScoreFilter = TraceThreadFilter.builder() - .field(TraceThreadField.FEEDBACK_SCORES) - .operator(operator) - .key(targetScoreName) - .value(filterValueFunction.apply(targetScoreValue)) - .build(); - - // When & Then - if (!stream) { - assertThreadPage(null, projectId, expectedThreads, List.of(feedbackScoreFilter), Map.of(), apiKey, - workspaceName); - } else { - assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, - List.of(feedbackScoreFilter)); - } - } - - private static boolean isExpected(Set expectedThreadIndices, UUID threadId, List allThreadIds) { - return expectedThreadIndices.stream() - .anyMatch(index -> allThreadIds.get(index).toString().equals(threadId.toString())); - } - } - - private List getExpectedThreads(List expectedTraces, UUID projectId, String threadId, - List spans, TraceThreadStatus status) { - return getExpectedThreads(expectedTraces, projectId, threadId, spans, status, null); - } - - private List getExpectedThreads(List expectedTraces, UUID projectId, String threadId, - List spans, TraceThreadStatus status, List feedbackScores) { - - return expectedTraces.isEmpty() - ? List.of() - : List.of(TraceThread.builder() - .firstMessage(expectedTraces.stream().min(Comparator.comparing(Trace::startTime)).orElseThrow() - .input()) - .lastMessage(expectedTraces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() - .output()) - .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision( - expectedTraces.stream().min(Comparator.comparing(Trace::startTime)).orElseThrow() - .startTime(), - expectedTraces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() - .endTime())) - .projectId(projectId) - .createdBy(USER) - .startTime(expectedTraces.stream().min(Comparator.comparing(Trace::startTime)).orElseThrow() - .startTime()) - .endTime(expectedTraces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() - .endTime()) - .numberOfMessages(expectedTraces.size() * 2L) - .id(threadId) - .totalEstimatedCost(Optional.ofNullable(getTotalEstimatedCost(spans)) - .filter(value -> value.compareTo(BigDecimal.ZERO) > 0) - .orElse(null)) - .usage(getUsage(spans)) - .status(status) - .feedbackScores(feedbackScores) - .createdAt(expectedTraces.stream().min(Comparator.comparing(Trace::createdAt)).orElseThrow() - .createdAt()) - .lastUpdatedAt( - expectedTraces.stream().max(Comparator.comparing(Trace::lastUpdatedAt)).orElseThrow() - .lastUpdatedAt()) - .build()); - } - - private Map getUsage(List spans) { - return Optional.ofNullable(spans) - .map(this::aggregateSpansUsage) - .filter(not(Map::isEmpty)) - .orElse(null); - } - - private BigDecimal getTotalEstimatedCost(List spans) { - boolean shouldUseTotalEstimatedCostField = spans.stream().allMatch(span -> span.totalEstimatedCost() != null); - - if (shouldUseTotalEstimatedCostField) { - return spans.stream() - .map(Span::totalEstimatedCost) - .reduce(BigDecimal.ZERO, BigDecimal::add); - } - - return calculateEstimatedCost(spans); - } - - @Nested - @DisplayName("Find traces:") - @TestInstance(TestInstance.Lifecycle.PER_CLASS) - class FindTraces { - - @ParameterizedTest - @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments") - void findWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) { - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = Stream.of(createTrace()) - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(null) - .input(original) - .output(original) - .metadata(original) - .build()) - .toList(); - - traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - - var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .queryParam("page", 1) - .queryParam("size", 5) - .queryParam("project_name", projectName) - .queryParam("truncate", truncate) - .request() - .header(HttpHeaders.AUTHORIZATION, API_KEY) - .header(WORKSPACE_HEADER, TEST_WORKSPACE) - .get(); - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); - - var actualPage = actualResponse.readEntity(Trace.TracePage.class); - var actualTraces = actualPage.content(); - - assertThat(actualTraces).hasSize(1); - - var expectedTraces = traces.stream() - .map(trace -> trace.toBuilder() - .input(expected) - .output(expected) - .metadata(expected) - .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(trace.startTime(), - trace.endTime())) - .build()) - .toList(); - - assertThat(actualTraces) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS_TRACES) - .containsExactlyElementsOf(expectedTraces); - } - - @ParameterizedTest - @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments") - void searchWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) { - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = Stream.of(createTrace()) - .map(trace -> trace.toBuilder() - .projectName(projectName) - .usage(null) - .input(original) - .output(original) - .metadata(original) - .build()) - .toList(); - - traceResourceClient.batchCreateTraces(traces, API_KEY, TEST_WORKSPACE); - - TraceSearchStreamRequest streamRequest = TraceSearchStreamRequest.builder() - .truncate(truncate) - .projectName(projectName) - .limit(5) - .build(); - - var actualTraces = traceResourceClient.getStreamAndAssertContent(API_KEY, TEST_WORKSPACE, streamRequest); - - assertThat(actualTraces).hasSize(1); - - var expectedTraces = traces.stream() - .map(trace -> trace.toBuilder() - .input(expected) - .output(expected) - .metadata(expected) - .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(trace.startTime(), - trace.endTime())) - .build()) - .toList(); - - TraceAssertions.assertTraces(actualTraces, expectedTraces, USER); - } - - @ParameterizedTest - @ValueSource(booleans = {true, false}) - void whenUsingPagination__thenReturnTracesPaginated(boolean stream) { - - String workspaceName = UUID.randomUUID().toString(); - String workspaceId = UUID.randomUUID().toString(); - String apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .comments(null) - .totalEstimatedCost(null) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var expectedTraces = traces.stream() - .sorted(Comparator.comparing(Trace::id).reversed()) - .toList(); - - int pageSize = 2; - - if (stream) { - AtomicReference lastId = new AtomicReference<>(null); - Lists.partition(expectedTraces, pageSize) - .forEach(trace -> { - var actualTraces = traceResourceClient.getStreamAndAssertContent(apiKey, workspaceName, - TraceSearchStreamRequest.builder() - .projectName(projectName) - .lastRetrievedId(lastId.get()) - .limit(pageSize) - .build()); - - TraceAssertions.assertTraces(actualTraces, trace, USER); - - lastId.set(actualTraces.getLast().id()); - }); - } else { - for (int i = 0; i < expectedTraces.size() / pageSize; i++) { - int page = i + 1; - getAndAssertPage( - page, - pageSize, - projectName, - null, - List.of(), - expectedTraces.subList(i * pageSize, Math.min((i + 1) * pageSize, expectedTraces.size())), - List.of(), - workspaceName, - apiKey, - List.of(), - traces.size(), Set.of()); - } - } - } - - @ParameterizedTest - @ValueSource(booleans = {true, false}) - void whenFilterByVisibilityScoreEqual__thenReturnTracesFiltered(boolean stream) { - - String workspaceName = UUID.randomUUID().toString(); - String workspaceId = UUID.randomUUID().toString(); - String apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .comments(null) - .totalEstimatedCost(null) - .build()) - .toList(); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - TraceFilter filter = TraceFilter.builder() - .field(TraceField.VISIBILITY_MODE) - .operator(Operator.EQUAL) - .value(VisibilityMode.DEFAULT.getValue()) - .build(); - - var actualTraces = traceResourceClient.getStreamAndAssertContent(apiKey, workspaceName, - TraceSearchStreamRequest.builder() - .projectName(projectName) - .filters(List.of(filter)) - .build()); - - if (stream) { - TraceAssertions.assertTraces(actualTraces, traces.reversed(), USER); - } else { - getAndAssertPage( - 1, - 100, - projectName, - null, - List.of(filter), - traces.reversed(), - List.of(), - workspaceName, - apiKey, - List.of(), - traces.size(), Set.of()); - } - } - - @ParameterizedTest - @MethodSource - void whenFilterByCustomFilter__thenReturnTracesFiltered(String key, String value, Operator operator) { - - String workspaceName = UUID.randomUUID().toString(); - String workspaceId = UUID.randomUUID().toString(); - String apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .comments(null) - .totalEstimatedCost(null) - .build()) - .collect(toCollection(ArrayList::new)); - - traces.set(0, traces.getFirst().toBuilder() - .input(JsonUtils - .getJsonNodeFromString("{\"model\":[{\"year\":2024,\"version\":\"OpenAI, " + - "Chat-GPT 4.0\",\"trueFlag\":true,\"nullField\":null}]}")) - .build()); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - TraceFilter filter = TraceFilter.builder() - .field(CUSTOM) - .operator(operator) - .key(key) - .value(value) - .build(); - - getAndAssertPage( - 1, - 100, - projectName, - null, - List.of(filter), - List.of(traces.getFirst()), - traces.subList(1, traces.size()), - workspaceName, - apiKey, - List.of(), - 1, Set.of()); - } - - private Stream whenFilterByCustomFilter__thenReturnTracesFiltered() { + private Stream getFeedbackScoreFilterTestArguments() { return Stream.of( Arguments.of( - "input.model[0].year", - "2024", - Operator.EQUAL), + true, + Operator.EQUAL, + generateExpectedIndices(), + (BiFunction>) this::generateExpectedEqualsMatch, + (BiFunction>) this::generateUnexpectedEqualsMatch, + (Function) BigDecimal::toString), // Filter value function for EQUAL Arguments.of( - "input.model[0].year", - "2025", - Operator.LESS_THAN), + false, + Operator.EQUAL, + generateExpectedIndices(), + (BiFunction>) this::generateExpectedEqualsMatch, + (BiFunction>) this::generateUnexpectedEqualsMatch, + (Function) BigDecimal::toString), Arguments.of( - "input", - "Chat-GPT 4.0", - Operator.CONTAINS)); - } - - @ParameterizedTest - @MethodSource - void getTracesByProject__whenSortingByValidFields__thenReturnTracesSorted(Comparator comparator, - SortingField sorting) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> { - var llmSpanCount = RandomUtils.secure().randomInt(1, 7); - return trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) - .comments(null) - .spanCount(llmSpanCount + RandomUtils.secure().randomInt(1, 7)) - .llmSpanCount(llmSpanCount) - .build(); - }) - .map(trace -> trace.toBuilder() - .duration(trace.startTime().until(trace.endTime(), ChronoUnit.MICROS) / 1000.0) - .build()) - .toList(); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - var spans = traces.stream() - .flatMap(trace -> IntStream.range(0, trace.spanCount()) - .mapToObj(i -> factory.manufacturePojo(Span.class).toBuilder() - .usage(Map.of("completion_tokens", RandomUtils.secure().randomInt())) - .projectName(projectName) - .traceId(trace.id()) - .type(i < trace.llmSpanCount() ? SpanType.llm : SpanType.general) - .build())) - .toList(); - - spanResourceClient.batchCreateSpans(spans, apiKey, workspaceName); - - var spansByTrace = spans.stream().collect(Collectors.groupingBy(Span::traceId)); - traces = traces.stream() - .map(t -> t.toBuilder() - .usage(aggregateSpansUsage(spansByTrace.get(t.id()))) - .build()) - .toList(); - - var expectedTraces = traces.stream() - .sorted(comparator) - .toList(); - - List sortingFields = List.of(sorting); - - getAndAssertPage(workspaceName, projectName, null, List.of(), traces, expectedTraces, List.of(), apiKey, - sortingFields, Set.of()); - } - - @Test - void createAndRetrieveTraces__spanCountReflectsActualSpans_andTotalCountMatches() { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - - // Create traces with varying spanCount values - List traces = IntStream.range(0, 5) - .mapToObj(i -> createTrace().toBuilder() - .projectId(null) - .projectName(projectName) - .spanCount(i * 3) // e.g., 0, 3, 6, 9, 12 - .usage(null) - .feedbackScores(null) - .endTime(Instant.now()) - .comments(null) - .build()) - .collect(Collectors.toList()); - - int expectedTotalSpanCount = traces.stream().mapToInt(Trace::spanCount).sum(); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); - - // For each trace, create the actual number of spans matching the spanCount - List allSpans = new ArrayList<>(); - for (Trace trace : traces) { - List spansForTrace = IntStream.range(0, trace.spanCount()) - .mapToObj(j -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(projectName) - .type(SpanType.llm) - .traceId(trace.id()) - .build()) - .toList(); - allSpans.addAll(spansForTrace); - } - spanResourceClient.batchCreateSpans(allSpans, apiKey, workspaceName); - - // Retrieve traces from the API - UUID projectId = getProjectId(projectName, workspaceName, apiKey); - Trace.TracePage resultPage = traceResourceClient.getTraces(projectName, projectId, apiKey, workspaceName, - List.of(), List.of(), 100, Map.of()); - List returnedTraces = resultPage.content(); - - // Check that all created traces are present and have the correct spanCount - for (Trace created : traces) { - returnedTraces.stream() - .filter(returned -> returned.id().equals(created.id())) - .findFirst() - .ifPresentOrElse(returned -> { - assertThat(returned.spanCount()) - .as("Trace with id %s should have spanCount %d", created.id(), created.spanCount()) - .isEqualTo(created.spanCount()); - assertThat(returned.llmSpanCount()) - .as("Trace with id %s should have llmSpanCount %d", created.id(), - created.spanCount()) - .isEqualTo(created.spanCount()); - }, - () -> assertThat(false) - .as("Trace with id %s should be present", created.id()) - .isTrue()); - } - - int actualTotalSpanCount = returnedTraces.stream() - .filter(rt -> traces.stream().anyMatch(t -> t.id().equals(rt.id()))) - .mapToInt(Trace::spanCount) - .sum(); - - assertThat(actualTotalSpanCount) - .as("Total spanCount across all traces should match the expected total") - .isEqualTo(expectedTotalSpanCount); - - int actualTotalLlmSpanCount = returnedTraces.stream() - .filter(rt -> traces.stream().anyMatch(t -> t.id().equals(rt.id()))) - .mapToInt(Trace::llmSpanCount) - .sum(); - - assertThat(actualTotalLlmSpanCount) - .as("Total llmSpanCount across all traces should match the expected total") - .isEqualTo(expectedTotalSpanCount); - } - - private Stream getTracesByProject__whenSortingByValidFields__thenReturnTracesSorted() { - - Comparator inputComparator = Comparator.comparing(trace -> trace.input().toString()); - Comparator outputComparator = Comparator.comparing(trace -> trace.output().toString()); - Comparator metadataComparator = Comparator.comparing(trace -> trace.metadata().toString()); - Comparator tagsComparator = Comparator.comparing(trace -> trace.tags().toString()); - Comparator errorInfoComparator = Comparator.comparing(trace -> trace.errorInfo().toString()); - Comparator usageComparator = Comparator.comparing(trace -> trace.usage().get("completion_tokens")); - - return Stream.of( - Arguments.of(Comparator.comparing(Trace::name), - SortingField.builder().field(SortableFields.NAME).direction(Direction.ASC).build()), - Arguments.of(Comparator.comparing(Trace::name).reversed(), - SortingField.builder().field(SortableFields.NAME).direction(Direction.DESC).build()), - Arguments.of(Comparator.comparing(Trace::startTime), - SortingField.builder().field(SortableFields.START_TIME).direction(Direction.ASC).build()), - Arguments.of(Comparator.comparing(Trace::startTime).reversed(), - SortingField.builder().field(SortableFields.START_TIME).direction(Direction.DESC).build()), - Arguments.of(Comparator.comparing(Trace::endTime), - SortingField.builder().field(SortableFields.END_TIME).direction(Direction.ASC).build()), - Arguments.of(Comparator.comparing(Trace::endTime).reversed(), - SortingField.builder().field(SortableFields.END_TIME).direction(Direction.DESC).build()), + true, + Operator.IS_NOT_EMPTY, + generateExpectedIndices(), + (BiFunction>) (name, + value) -> generateNotEmptyMatch(name), + (BiFunction>) (name, + value) -> generateUnexpectedNotEmptyMatch(name), + (Function) value -> ""), // Empty value for IS_NOT_EMPTY Arguments.of( - Comparator.comparing(Trace::duration) - .thenComparing(Comparator.comparing(Trace::id).reversed()), - SortingField.builder().field(SortableFields.DURATION).direction(Direction.ASC).build()), + false, + Operator.IS_NOT_EMPTY, + generateExpectedIndices(), + (BiFunction>) (name, + value) -> generateNotEmptyMatch(name), + (BiFunction>) (name, + value) -> generateUnexpectedNotEmptyMatch(name), + (Function) value -> ""), Arguments.of( - Comparator.comparing(Trace::duration).reversed() - .thenComparing(Comparator.comparing(Trace::id).reversed()), - SortingField.builder().field(SortableFields.DURATION).direction(Direction.DESC).build()), - Arguments.of(inputComparator, - SortingField.builder().field(SortableFields.INPUT).direction(Direction.ASC).build()), - Arguments.of(inputComparator.reversed(), - SortingField.builder().field(SortableFields.INPUT).direction(Direction.DESC).build()), - Arguments.of(outputComparator, - SortingField.builder().field(SortableFields.OUTPUT).direction(Direction.ASC).build()), - Arguments.of(outputComparator.reversed(), - SortingField.builder().field(SortableFields.OUTPUT).direction(Direction.DESC).build()), - Arguments.of(metadataComparator, - SortingField.builder().field(SortableFields.METADATA).direction(Direction.ASC).build()), - Arguments.of(metadataComparator.reversed(), - SortingField.builder().field(SortableFields.METADATA).direction(Direction.DESC).build()), - Arguments.of(tagsComparator, - SortingField.builder().field(SortableFields.TAGS).direction(Direction.ASC).build()), - Arguments.of(tagsComparator.reversed(), - SortingField.builder().field(SortableFields.TAGS).direction(Direction.DESC).build()), - Arguments.of(Comparator.comparing(Trace::id), - SortingField.builder().field(SortableFields.ID).direction(Direction.ASC).build()), - Arguments.of(Comparator.comparing(Trace::id).reversed(), - SortingField.builder().field(SortableFields.ID).direction(Direction.DESC).build()), - Arguments.of(errorInfoComparator, - SortingField.builder().field(SortableFields.ERROR_INFO).direction(Direction.ASC).build()), - Arguments.of(errorInfoComparator.reversed(), - SortingField.builder().field(SortableFields.ERROR_INFO).direction(Direction.DESC).build()), - Arguments.of(Comparator.comparing(Trace::threadId), SortingField.builder() - .field(SortableFields.THREAD_ID).direction(Direction.ASC).build()), - Arguments.of(Comparator.comparing(Trace::threadId).reversed(), SortingField.builder() - .field(SortableFields.THREAD_ID).direction(Direction.DESC).build()), - Arguments.of(Comparator.comparing(Trace::spanCount) - .thenComparing(Comparator.comparing(Trace::id).reversed()), - SortingField.builder().field(SortableFields.SPAN_COUNT).direction(Direction.ASC).build()), - Arguments.of(Comparator.comparing(Trace::spanCount).reversed() - .thenComparing(Comparator.comparing(Trace::id).reversed()), - SortingField.builder().field(SortableFields.SPAN_COUNT).direction(Direction.DESC).build()), - Arguments.of(Comparator.comparing(Trace::llmSpanCount) - .thenComparing(Comparator.comparing(Trace::id).reversed()), - SortingField.builder().field(SortableFields.LLM_SPAN_COUNT).direction(Direction.ASC) - .build()), - Arguments.of(Comparator.comparing(Trace::llmSpanCount).reversed() - .thenComparing(Comparator.comparing(Trace::id).reversed()), - SortingField.builder().field(SortableFields.LLM_SPAN_COUNT).direction(Direction.DESC) - .build()), - Arguments.of(usageComparator, - SortingField.builder().field("usage.completion_tokens").direction(Direction.ASC).build()), - Arguments.of(usageComparator.reversed(), - SortingField.builder().field("usage.completion_tokens").direction(Direction.DESC).build())); - } - - @Test - void getTracesByProject__whenSortingByInvalidField__thenReturn400() { - var field = RandomStringUtils.secure().nextAlphanumeric(10); - var expectedError = new io.dropwizard.jersey.errors.ErrorMessage( - 400, - "Invalid sorting fields '%s'".formatted(field)); - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - - var sortingFields = List.of(SortingField.builder().field(field).direction(Direction.ASC).build()); - var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .queryParam("project_name", projectName) - .queryParam("sorting", - URLEncoder.encode(JsonUtils.writeValueAsString(sortingFields), StandardCharsets.UTF_8)) - .request() - .header(HttpHeaders.AUTHORIZATION, API_KEY) - .header(WORKSPACE_HEADER, TEST_WORKSPACE) - .get(); - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + true, + Operator.IS_EMPTY, + generateExpectedIndices(), + (BiFunction>) (name, + value) -> generateIsEmptyMatch(name), + (BiFunction>) (name, + value) -> generateNotEmptyMatch(name), + (Function) value -> ""), // Empty value for IS_EMPTY + Arguments.of( + false, + Operator.IS_EMPTY, + generateExpectedIndices(), + (BiFunction>) (name, + value) -> generateIsEmptyMatch(name), + (BiFunction>) (name, + value) -> generateNotEmptyMatch(name), + (Function) value -> "") - var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); - assertThat(actualError).isEqualTo(expectedError); + ); } - @ParameterizedTest - @EnumSource(Direction.class) - void getTracesByProject__whenSortingByFeedbackScores__thenReturnTracesSorted(Direction direction) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + private Set generateExpectedIndices() { + return new HashSet<>(List.of(RandomUtils.secure().randomInt(0, 5), RandomUtils.secure().randomInt(0, 5))); + } - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) + private List generateIsEmptyMatch(String name) { + return PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItemThread.class) .stream() - .map(trace -> trace.toBuilder() - .projectId(null) - .projectName(projectName) - .usage(null) - .feedbackScores(null) - .endTime(trace.startTime().plus(randomNumber(), ChronoUnit.MILLIS)) - .comments(null) - .build()) - .map(trace -> trace.toBuilder() - .duration(trace.startTime().until(trace.endTime(), ChronoUnit.MICROS) / 1000.0) - .build()) - .collect(Collectors.toCollection(ArrayList::new)); - - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + .filter(score -> !score.name().equals(name)) + .toList(); + } - List scoreForTrace = PodamFactoryUtils.manufacturePojoList(factory, - FeedbackScoreBatchItem.class); + private List generateNotEmptyMatch(String name) { + List scores = PodamFactoryUtils.manufacturePojoList(factory, + FeedbackScoreBatchItemThread.class); - List allScores = new ArrayList<>(); - for (Trace trace : traces) { - for (FeedbackScoreBatchItem item : scoreForTrace) { + scores.set(0, scores.getFirst().toBuilder() + .name(name) + .build()); - if (traces.getLast().equals(trace) && scoreForTrace.getFirst().equals(item)) { - continue; - } + return scores; + } - allScores.add(item.toBuilder() - .id(trace.id()) - .projectName(trace.projectName()) - .value(factory.manufacturePojo(BigDecimal.class).abs()) - .build()); - } - } + private List generateUnexpectedNotEmptyMatch(String name) { + return PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItemThread.class) + .stream() + .filter(score -> !score.name().equals(name)) + .toList(); + } - traceResourceClient.feedbackScores(allScores, apiKey, workspaceName); + private List generateExpectedEqualsMatch(String name, BigDecimal value) { + List scores = PodamFactoryUtils.manufacturePojoList(factory, + FeedbackScoreBatchItemThread.class); - var sortingField = new SortingField( - "feedback_scores.%s".formatted(scoreForTrace.getFirst().name()), - direction); + scores.set(0, scores.getFirst().toBuilder() + .name(name) + .value(value) + .build()); - Comparator comparing = Comparator.comparing( - (Trace trace) -> trace.feedbackScores() - .stream() - .filter(score -> score.name().equals(scoreForTrace.getFirst().name())) - .findFirst() - .map(FeedbackScore::value) - .orElse(null), - direction == Direction.ASC - ? Comparator.nullsFirst(Comparator.naturalOrder()) - : Comparator.nullsLast(Comparator.reverseOrder())) - .thenComparing(Comparator.comparing(Trace::id).reversed()); + return scores; + } - var expectedTraces = traces.stream() - .map(trace -> trace.toBuilder() - .feedbackScores(allScores - .stream() - .filter(score -> score.id().equals(trace.id())) - .map(scores -> FeedbackScore.builder() - .name(scores.name()) - .value(scores.value()) - .categoryName(scores.categoryName()) - .source(scores.source()) - .reason(scores.reason()) - .build()) - .toList()) - .build()) - .sorted(comparing) - .toList(); + private List generateUnexpectedEqualsMatch(String name, BigDecimal value) { + List scores = PodamFactoryUtils.manufacturePojoList(factory, + FeedbackScoreBatchItemThread.class); - List sortingFields = List.of(sortingField); + scores.set(0, scores.getFirst().toBuilder() + .name(name) + .build()); - getAndAssertPage(workspaceName, projectName, null, List.of(), traces, expectedTraces, List.of(), apiKey, - sortingFields, Set.of()); + return scores; } @ParameterizedTest - @EnumSource(Trace.TraceField.class) - void getTracesByProject__whenExcludeParamIdDefined__thenReturnSpanExcludingFields(Trace.TraceField field) { + @MethodSource("getFeedbackScoreFilterTestArguments") + @DisplayName("When filtering by feedback score with different operators, should return matching threads") + void whenFilterByFeedbackScore__thenReturnThreadsWithMatchingFeedbackScore( + boolean stream, + Operator operator, + Set expectedThreadIndices, + BiFunction> matchingScoreFunction, + BiFunction> unmatchingScoreFunction, + Function filterValueFunction) { + + // Given var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); var workspaceId = UUID.randomUUID().toString(); var apiKey = UUID.randomUUID().toString(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class) - .stream() - .map(trace -> trace.toBuilder().projectName(projectName).build()) - .toList(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(10); + var project = factory.manufacturePojo(Project.class).toBuilder() + .name(projectName) + .build(); - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + UUID projectId = projectResourceClient.createProject(project, apiKey, workspaceName); - Map expectedComments = traces - .stream() - .map(trace -> Map.entry(trace.id(), - traceResourceClient.generateAndCreateComment(trace.id(), apiKey, workspaceName, 201))) - .collect(toMap(Map.Entry::getKey, Map.Entry::getValue)); + // Create threads with different feedback scores + var allThreadIds = PodamFactoryUtils.manufacturePojoList(factory, UUID.class); - traces = traces.stream() - .map(span -> span.toBuilder() - .comments(List.of(expectedComments.get(span.id()))) - .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(span.startTime(), - span.endTime())) + // Create traces for threads + var allTraces = allThreadIds + .stream() + .map(threadId -> createTrace().toBuilder() + .threadId(threadId.toString()) + .projectId(null) + .projectName(projectName) + .startTime(Instant.now().minusSeconds(3).truncatedTo(ChronoUnit.MICROS)) + .lastUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MICROS)) .build()) .toList(); - List spans = traces.stream() - .map(trace -> factory.manufacturePojo(Span.class).toBuilder() - .projectName(trace.projectName()) - .traceId(trace.id()) - .build()) - .toList(); + allTraces.forEach(trace -> traceResourceClient.createTrace(trace, apiKey, workspaceName)); - batchCreateSpansAndAssert(spans, apiKey, workspaceName); + // Add feedback scores with different values + Map threadIdAndLastUpdateAts = new HashMap<>(); - traces = traces.stream() - .map(trace -> trace.toBuilder() - .totalEstimatedCost(spans.stream() - .filter(span -> span.traceId().equals(trace.id())) - .map(Span::totalEstimatedCost) - .reduce(BigDecimal.ZERO, BigDecimal::add)) - .spanCount((int) spans.stream() - .filter(span -> span.traceId().equals(trace.id())) - .count()) - .usage(spans.stream() - .filter(span -> span.traceId().equals(trace.id())) - .map(Span::usage) - .flatMap(map -> map.entrySet().stream()) - .collect(Collectors.groupingBy(Map.Entry::getKey, - Collectors.summingLong(Map.Entry::getValue)))) - .build()) - .toList(); + Mono.delay(Duration.ofMillis(500)).block(); + allThreadIds.forEach(threadId -> { + threadIdAndLastUpdateAts.put(threadId.toString(), Instant.now()); + traceResourceClient.closeTraceThread(threadId.toString(), null, projectName, apiKey, workspaceName); + }); - List finalTraces = traces; - List scoreForSpan = IntStream.range(0, traces.size()) - .mapToObj(i -> initFeedbackScoreItem() - .projectName(finalTraces.get(i).projectName()) - .id(finalTraces.get(i).id()) - .build()) - .collect(Collectors.toList()); + String targetScoreName = RandomStringUtils.secure().nextAlphanumeric(30); + BigDecimal targetScoreValue = factory.manufacturePojo(BigDecimal.class); - traceResourceClient.feedbackScores(scoreForSpan, apiKey, workspaceName); + // Create feedback scores based on the provided map + List expectedScores = allThreadIds + .stream() + .filter(threadId -> isExpected(expectedThreadIndices, threadId, allThreadIds)) + .flatMap(threadId -> { + return matchingScoreFunction.apply(targetScoreName, targetScoreValue).stream() + .map(item -> item.toBuilder() + .threadId(threadId.toString()) + .projectName(projectName) + .build()); + }).collect(Collectors.toList()); - traces = traces.stream() - .map(trace -> trace.toBuilder() - .feedbackScores( - scoreForSpan - .stream() - .filter(score -> score.id().equals(trace.id())) - .map(scores -> FeedbackScore.builder() - .name(scores.name()) - .value(scores.value()) - .categoryName(scores.categoryName()) - .source(scores.source()) - .reason(scores.reason()) - .build()) - .toList()) - .build()) - .toList(); + List unexpectedScores = allThreadIds.stream() + .filter(threadId -> !isExpected(expectedThreadIndices, threadId, allThreadIds)) + .flatMap(threadId -> { + return unmatchingScoreFunction.apply(targetScoreName, targetScoreValue).stream() + .map(item -> item.toBuilder() + .threadId(threadId.toString()) + .projectName(projectName) + .build()); + }).collect(Collectors.toList()); - List guardrailsByTraceId = traces.stream() - .map(trace -> guardrailsGenerator.generateGuardrailsForTrace(trace.id(), randomUUID(), - trace.projectName())) - .flatMap(Collection::stream) + List scoreItems = Stream + .concat(expectedScores.stream(), unexpectedScores.stream()) .toList(); - traces = traces.stream() - .map(trace -> trace.toBuilder() - .guardrailsValidations( - GuardrailsMapper.INSTANCE.mapToValidations( - guardrailsByTraceId - .stream() - .filter(gr -> gr.entityId().equals(trace.id())) - .toList())) - .build()) + // Create feedback scores for threads + Instant feedbackScoreCreationTime = Instant.now(); + traceResourceClient.threadFeedbackScores(scoreItems, apiKey, workspaceName); + + // Determine expected threads based on indices + var expectedThreadIds = allThreadIds.reversed() + .stream() + .filter(threadId -> isExpected(expectedThreadIndices, threadId, allThreadIds)) + .map(UUID::toString) .toList(); - guardrailsResourceClient.addBatch(guardrailsByTraceId, apiKey, workspaceName); + // Create expected threads with ALL feedback scores from matching threads + Comparator comparing = Comparator + .comparing((TraceThread traceThread) -> threadIdAndLastUpdateAts.get(traceThread.id())).reversed(); - traces = traces.stream() - .map(span -> EXCLUDE_FUNCTIONS.get(field).apply(span)) - .toList(); + List expectedThreads = expectedThreadIds.stream() + .map(threadId -> { + // Get ALL feedback scores for this thread (both expected and unexpected) + var allFeedbackScoresForThread = scoreItems.stream() + .filter(item -> item.threadId().equals(threadId)) + .map(item -> createExpectedFeedbackScore(item, feedbackScoreCreationTime)) + .toList(); + + var traces = allTraces.stream() + .filter(trace -> trace.threadId().equals(threadId)) + .toList(); - Set exclude = Set.of(field); + return getExpectedThreads(traces, projectId, threadId, List.of(), + TraceThreadStatus.INACTIVE, allFeedbackScoresForThread).getFirst(); + }) + .sorted(comparing) + .toList(); - getAndAssertPage(workspaceName, projectName, null, List.of(), traces, traces.reversed(), List.of(), apiKey, - List.of(), exclude); + // Create filter for the specific feedback score + var feedbackScoreFilter = TraceThreadFilter.builder() + .field(TraceThreadField.FEEDBACK_SCORES) + .operator(operator) + .key(targetScoreName) + .value(filterValueFunction.apply(targetScoreValue)) + .build(); + // When & Then + if (!stream) { + assertThreadPage(null, projectId, expectedThreads, List.of(feedbackScoreFilter), Map.of(), apiKey, + workspaceName); + } else { + assertTheadStream(null, projectId, apiKey, workspaceName, expectedThreads, + List.of(feedbackScoreFilter)); + } } - @Test - @DisplayName("should handle filter with percent characters in value correctly") - void shouldHandleTracesWithPercentCharactersInName() { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); + private static boolean isExpected(Set expectedThreadIndices, UUID threadId, List allThreadIds) { + return expectedThreadIndices.stream() + .anyMatch(index -> allThreadIds.get(index).toString().equals(threadId.toString())); + } + } - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + private List getExpectedThreads(List expectedTraces, UUID projectId, String threadId, + List spans, TraceThreadStatus status) { + return getExpectedThreads(expectedTraces, projectId, threadId, spans, status, null); + } - var projectName = RandomStringUtils.secure().nextAlphanumeric(10); - var traceName = "test%"; + private List getExpectedThreads(List expectedTraces, UUID projectId, String threadId, + List spans, TraceThreadStatus status, List feedbackScores) { - // Create a trace with % characters in the name - var traces = List.of(createTrace().toBuilder() - .projectName(projectName) - .name(traceName) - .usage(null) - .feedbackScores(null) - .threadId(null) - .comments(null) - .totalEstimatedCost(null) - .build()); + return expectedTraces.isEmpty() + ? List.of() + : List.of(TraceThread.builder() + .firstMessage(expectedTraces.stream().min(Comparator.comparing(Trace::startTime)).orElseThrow() + .input()) + .lastMessage(expectedTraces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() + .output()) + .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision( + expectedTraces.stream().min(Comparator.comparing(Trace::startTime)).orElseThrow() + .startTime(), + expectedTraces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() + .endTime())) + .projectId(projectId) + .createdBy(USER) + .startTime(expectedTraces.stream().min(Comparator.comparing(Trace::startTime)).orElseThrow() + .startTime()) + .endTime(expectedTraces.stream().max(Comparator.comparing(Trace::endTime)).orElseThrow() + .endTime()) + .numberOfMessages(expectedTraces.size() * 2L) + .id(threadId) + .totalEstimatedCost(Optional.ofNullable(getTotalEstimatedCost(spans)) + .filter(value -> value.compareTo(BigDecimal.ZERO) > 0) + .orElse(null)) + .usage(getUsage(spans)) + .status(status) + .feedbackScores(feedbackScores) + .createdAt(expectedTraces.stream().min(Comparator.comparing(Trace::createdAt)).orElseThrow() + .createdAt()) + .lastUpdatedAt( + expectedTraces.stream().max(Comparator.comparing(Trace::lastUpdatedAt)).orElseThrow() + .lastUpdatedAt()) + .build()); + } - traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName); + private Map getUsage(List spans) { + return Optional.ofNullable(spans) + .map(this::aggregateSpansUsage) + .filter(not(Map::isEmpty)) + .orElse(null); + } - // Create a filter to search for the trace by name - var filter = TraceFilter.builder() - .field(TraceField.NAME) - .operator(Operator.EQUAL) - .value(traceName) - .build(); + private BigDecimal getTotalEstimatedCost(List spans) { + boolean shouldUseTotalEstimatedCostField = spans.stream().allMatch(span -> span.totalEstimatedCost() != null); - getAndAssertPage(workspaceName, projectName, null, List.of(filter), traces, traces, List.of(), - apiKey, null, Set.of()); + if (shouldUseTotalEstimatedCostField) { + return spans.stream() + .map(Span::totalEstimatedCost) + .reduce(BigDecimal.ZERO, BigDecimal::add); } - } - private Integer randomNumber() { - return randomNumber(10, 99); + return calculateEstimatedCost(spans); } private static int randomNumber(int minValue, int maxValue) { @@ -6632,19 +2304,6 @@ private void getAndAssertPageSpans( } } - private List updateFeedbackScore(List feedbackScores, int index, double val) { - feedbackScores.set(index, feedbackScores.get(index).toBuilder() - .value(BigDecimal.valueOf(val)) - .build()); - return feedbackScores; - } - - private List updateFeedbackScore( - List destination, List source, int index) { - destination.set(index, source.get(index).toBuilder().build()); - return destination; - } - @Nested @DisplayName("Thread Feedback Scores Creation:") @TestInstance(TestInstance.Lifecycle.PER_CLASS) @@ -11016,15 +6675,6 @@ private Map aggregateSpansUsage(List spans) { .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, Long::sum)); } - private Trace updateSpanCounts(Trace trace, List spans) { - return updateSpanCounts(List.of(trace), spans).getFirst(); - } - - private List updateSpanCounts(List traces, List spans) { - var spansByTraceId = spans.stream().collect(Collectors.groupingBy(Span::traceId)); - return updateSpanCounts(traces, spansByTraceId); - } - private List updateSpanCounts(List traces, Map> spansByTraceId) { return traces.stream() .map(trace -> { From de89734fd2cec79633fbd7c553007b44b8c8ed42 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Mon, 3 Nov 2025 17:36:55 +0100 Subject: [PATCH 4/9] [OPIK-2856] Address PR review comments: improve exception handling and type safety - Fix InstantParamConverter to catch specific DateTimeParseException instead of generic Exception - Add debug logging when falling back to epoch milliseconds parsing - Refactor anonymous ParamConverter class to named InstantConverter inner class for clarity - Suppress unchecked cast with @SuppressWarnings annotation - Fix MySQLContainerUtils return type to use MySQLContainer for type safety --- .../web/InstantParamConverter.java | 57 ++++++++++--------- 1 file changed, 29 insertions(+), 28 deletions(-) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java index 32846054474..4c9c9fa3e1f 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java @@ -9,6 +9,7 @@ import java.lang.annotation.Annotation; import java.lang.reflect.Type; import java.time.Instant; +import java.time.format.DateTimeParseException; /** * JAX-RS ParamConverter for automatic Instant conversion from query parameters. @@ -24,38 +25,38 @@ public ParamConverter getConverter(Class rawType, Type genericType, An return null; } - return (ParamConverter) new ParamConverter() { - @Override - public Instant fromString(String value) { - if (value == null || value.isEmpty()) { - return null; - } + return (ParamConverter) new InstantConverter(); + } + + private static class InstantConverter implements ParamConverter { + + @Override + public Instant fromString(String value) { + if (value == null || value.isEmpty()) { + return null; + } + try { + // Try parsing as ISO-8601 format first + return Instant.parse(value); + } catch (DateTimeParseException exception) { + log.debug("Failed to parse '{}' as ISO-8601, attempting to parse as milliseconds since epoch", value, + exception); try { - // Try parsing as ISO-8601 format first - return Instant.parse(value); - } catch (Exception exception) { - log.debug("Failed to parse '{}' as ISO-8601, attempting to parse as milliseconds since epoch", - value); - try { - // Fall back to parsing as milliseconds since epoch - long epochMillis = Long.parseLong(value); - return Instant.ofEpochMilli(epochMillis); - } catch (NumberFormatException numberFormatException) { - log.error( - "Invalid instant format: '{}'. Expected ISO-8601 (e.g., 2024-01-01T00:00:00Z) or milliseconds since epoch", - value); - throw new BadRequestException( - "Invalid instant format: '%s'. Expected ISO-8601 (e.g., 2024-01-01T00:00:00Z) or milliseconds since epoch" - .formatted(value)); - } + // Fall back to parsing as milliseconds since epoch + long epochMillis = Long.parseLong(value); + return Instant.ofEpochMilli(epochMillis); + } catch (NumberFormatException numberFormatException) { + throw new BadRequestException( + "Invalid instant format: '%s'. Expected ISO-8601 (e.g., 2024-01-01T00:00:00Z) or milliseconds since epoch" + .formatted(value)); } } + } - @Override - public String toString(Instant value) { - return value != null ? value.toString() : null; - } - }; + @Override + public String toString(Instant value) { + return value != null ? value.toString() : null; + } } } From b33aca0f19728f2ed587b2c3e96e32e0555b71f2 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Mon, 3 Nov 2025 17:47:34 +0100 Subject: [PATCH 5/9] [OPIK-2856] Fix InstantToUUIDMapperTest to match implementation - Update tests to reflect that toUpperBound uses next millisecond (+1ms) for inclusive BETWEEN queries - Remove outdated assertions expecting same timestamp in both bounds - Verify upper bound is lexicographically greater than lower bound - All 13 tests now passing --- .../opik/api/InstantToUUIDMapperTest.java | 74 ++++++++++--------- 1 file changed, 38 insertions(+), 36 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java index f3c83f5f6df..e9387152d90 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java @@ -71,9 +71,13 @@ void shouldEncodeSameTimestampInBothBounds() { String lowerTimestampPart = lowerBoundStr.substring(0, 12); String upperTimestampPart = upperBoundStr.substring(0, 12); - assertThat(lowerTimestampPart) - .as("Both bounds should encode the same timestamp") - .isEqualTo(upperTimestampPart); + // Upper bound uses next millisecond (+1ms), so timestamps should differ by 1 + UUID nextMillisLowerBound = InstantToUUIDMapper.toLowerBound(timestamp.plusMillis(1)); + String nextMillisTimestampPart = nextMillisLowerBound.toString().replace("-", "").substring(0, 12); + + assertThat(upperTimestampPart) + .as("Upper bound should encode the next millisecond") + .isEqualTo(nextMillisTimestampPart); } @Test @@ -86,16 +90,12 @@ void shouldHaveDifferentRandomPortions() { UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); // Then - // Extract the random portion (after the timestamp and version) - String lowerBoundStr = lowerBound.toString().replace("-", ""); - String upperBoundStr = upperBound.toString().replace("-", ""); - - String lowerRandomPart = lowerBoundStr.substring(13); // Skip timestamp+version - String upperRandomPart = upperBoundStr.substring(13); - - assertThat(lowerRandomPart) - .as("Random portions should differ (lower should be 0000, upper should be FFFF)") - .isNotEqualTo(upperRandomPart); + // Lower bound uses same timestamp with zero random bytes + // Upper bound uses next millisecond with zero random bytes + // Since they have different timestamps, they should be different UUIDs + assertThat(lowerBound) + .as("Upper and lower bounds should be different UUIDs") + .isNotEqualTo(upperBound); } @Test @@ -173,13 +173,17 @@ void shouldUseFFBytesForUpperBound() { Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); - UUID referenceUUID = OpenTelemetryMapper.convertOtelIdToUUIDv7( - new byte[]{-1, -1, -1, -1, -1, -1, -1, -1}, - timestamp.toEpochMilli()); // Then - assertThat(upperBound).isEqualTo(referenceUUID); + assertThat(lowerBound).isNotNull(); + assertThat(upperBound).isNotNull(); + + // Upper bound should be greater than lower bound + assertThat(upperBound.toString().compareTo(lowerBound.toString())) + .as("Upper bound should be lexicographically greater than lower bound") + .isGreaterThan(0); } @Test @@ -201,24 +205,22 @@ void shouldWorkWithBoundaryInstants() { @Test void shouldAllUUIDsWithSameTimestampShareTimestampPortion() { // Given - Instant queryTime = Instant.parse("2025-01-15T10:30:00Z"); - UUID lowerBound = InstantToUUIDMapper.toLowerBound(queryTime); - UUID upperBound = InstantToUUIDMapper.toUpperBound(queryTime); - UUID randomUUID = OpenTelemetryMapper.convertOtelIdToUUIDv7( - new byte[]{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}, - queryTime.toEpochMilli()); - - // Then - All UUIDs with the same timestamp should have the same timestamp portion - String lowerStr = lowerBound.toString().replace("-", ""); - String upperStr = upperBound.toString().replace("-", ""); - String randomStr = randomUUID.toString().replace("-", ""); - - // Extract timestamp (first 12 hex chars = 48 bits) - String lowerTimestamp = lowerStr.substring(0, 12); - String upperTimestamp = upperStr.substring(0, 12); - String randomTimestamp = randomStr.substring(0, 12); - - assertThat(lowerTimestamp).isEqualTo(upperTimestamp); - assertThat(lowerTimestamp).isEqualTo(randomTimestamp); + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + + // When + UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + UUID nextMillisLowerBound = InstantToUUIDMapper.toLowerBound(timestamp.plusMillis(1)); + + // Then + // Extract the timestamp portion (first 48 bits / first 12 hex chars) + String lowerBoundStr = lowerBound.toString().replace("-", ""); + String nextMillisStr = nextMillisLowerBound.toString().replace("-", ""); + + String lowerTimestampPart = lowerBoundStr.substring(0, 12); + String nextMillisTimestampPart = nextMillisStr.substring(0, 12); + + assertThat(nextMillisTimestampPart) + .as("Next millisecond should have different timestamp portion") + .isNotEqualTo(lowerTimestampPart); } } From fc34b96a4f110b84cf96c8df90c54d8d21cc29e0 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Tue, 4 Nov 2025 10:42:16 +0100 Subject: [PATCH 6/9] Remove setup duplicated code --- .../api/resources/utils/TestWorkspace.java | 17 ++ .../priv/GetTracesByProjectResourceTest.java | 269 ++++++------------ 2 files changed, 106 insertions(+), 180 deletions(-) create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestWorkspace.java diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestWorkspace.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestWorkspace.java new file mode 100644 index 00000000000..1e6b3f2225f --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestWorkspace.java @@ -0,0 +1,17 @@ +package com.comet.opik.api.resources.utils; + +/** + * Record for holding test workspace configuration data used across test classes. + * Encapsulates workspace-related test data like workspace name, ID, API key, and project name. + * + * @param workspaceName the name of the test workspace + * @param workspaceId the unique identifier of the test workspace + * @param apiKey the API key for the test workspace + * @param projectName the name of the test project + */ +public record TestWorkspace( + String workspaceName, + String workspaceId, + String apiKey, + String projectName) { +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java index d61283e991b..55d031d869d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/GetTracesByProjectResourceTest.java @@ -27,6 +27,7 @@ import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; import com.comet.opik.api.resources.utils.TestUtils; +import com.comet.opik.api.resources.utils.TestWorkspace; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.api.resources.utils.resources.AnnotationQueuesResourceClient; import com.comet.opik.api.resources.utils.resources.GuardrailsGenerator; @@ -4687,65 +4688,34 @@ private Stream provideBoundaryScenarios() { @MethodSource("provideBoundaryScenarios") void whenTimeParametersProvided_thenIncludeTracesWithinBounds( String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var workspace = setupTestWorkspace(); Instant baseTime = Instant.now(); Instant lowerBound = baseTime.minus(Duration.ofMinutes(10)); Instant upperBound = baseTime; // Create traces with UUIDs at specific boundary times - List allTraces = new ArrayList<>(); - - // Index 0: At exact lower bound (should be included) - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(lowerBound)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); + List allTraces = List.of( + createTraceAtTimestamp(workspace.projectName(), lowerBound, + "At exact lower bound (should be included)"), + createTraceAtTimestamp(workspace.projectName(), upperBound, + "At exact upper bound (should be included)"), + createTraceAtTimestamp(workspace.projectName(), lowerBound.plus(Duration.ofMinutes(5)), + "Between bounds (should be included)")); - // Index 1: At exact upper bound (should be included) - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(upperBound)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - // Index 2: Between bounds (should be included) - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(lowerBound.plus(Duration.ofMinutes(5)))) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - traceResourceClient.batchCreateTraces(allTraces, apiKey, workspaceName); + traceResourceClient.batchCreateTraces(allTraces, workspace.apiKey(), workspace.workspaceName()); var queryParams = Map.of( "from_time", lowerBound.toString(), "to_time", upperBound.toString()); // Clear projectName from traces since API returns projectName=null - allTraces = allTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - - // Sort by ID descending to match API response order - var expectedTraces = allTraces.stream() - .sorted(Comparator.comparing(Trace::id).reversed()) - .toList(); - + allTraces = normalizeTraces(allTraces); + var expectedTraces = sortByIdDescending(allTraces); var values = testAssertion.transformTestParams(allTraces, expectedTraces, List.of()); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), - values.unexpected(), values.all(), List.of(), queryParams); + testAssertion.assertTest(workspace.projectName(), null, workspace.apiKey(), workspace.workspaceName(), + values.expected(), values.unexpected(), values.all(), List.of(), queryParams); } @ParameterizedTest @@ -4753,62 +4723,24 @@ void whenTimeParametersProvided_thenIncludeTracesWithinBounds( @MethodSource("provideBoundaryScenarios") void whenTimeParametersProvided_thenExcludeTracesOutsideBounds( String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var workspace = setupTestWorkspace(); Instant baseTime = Instant.now(); Instant lowerBound = baseTime.minus(Duration.ofMinutes(10)); Instant upperBound = baseTime; // Create traces: 3 within bounds, 2 outside bounds - List allTraces = new ArrayList<>(); - - // Within bounds - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(lowerBound)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(upperBound)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(lowerBound.plus(Duration.ofMinutes(1)))) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - // Outside bounds (before lower) - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(lowerBound.minus(Duration.ofMinutes(1)))) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - // Outside bounds (after upper) - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(upperBound.plus(Duration.ofMinutes(1)))) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - traceResourceClient.batchCreateTraces(allTraces, apiKey, workspaceName); + List allTraces = List.of( + createTraceAtTimestamp(workspace.projectName(), lowerBound, "Within bounds"), + createTraceAtTimestamp(workspace.projectName(), upperBound, "Within bounds"), + createTraceAtTimestamp(workspace.projectName(), lowerBound.plus(Duration.ofMinutes(1)), + "Within bounds"), + createTraceAtTimestamp(workspace.projectName(), lowerBound.minus(Duration.ofMinutes(1)), + "Outside bounds (before lower)"), + createTraceAtTimestamp(workspace.projectName(), upperBound.plus(Duration.ofMinutes(1)), + "Outside bounds (after upper)")); + + traceResourceClient.batchCreateTraces(allTraces, workspace.apiKey(), workspace.workspaceName()); // Expected: indices 0, 1, 2 (within bounds) // Unexpected: indices 3, 4 (outside bounds) @@ -4819,20 +4751,15 @@ void whenTimeParametersProvided_thenExcludeTracesOutsideBounds( "from_time", lowerBound.toString(), "to_time", upperBound.toString()); - // Clear projectName from traces since API returns projectName=null - allTraces = allTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - expectedTraces = expectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - unexpectedTraces = unexpectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - - // Sort expectedTraces by ID descending to match API response order - expectedTraces = expectedTraces.stream() - .sorted(Comparator.comparing(Trace::id).reversed()) - .toList(); + allTraces = normalizeTraces(allTraces); + expectedTraces = normalizeTraces(expectedTraces); + unexpectedTraces = normalizeTraces(unexpectedTraces); + expectedTraces = sortByIdDescending(expectedTraces); var values = testAssertion.transformTestParams(allTraces, expectedTraces, unexpectedTraces); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), - values.unexpected(), values.all(), List.of(), queryParams); + testAssertion.assertTest(workspace.projectName(), null, workspace.apiKey(), workspace.workspaceName(), + values.expected(), values.unexpected(), values.all(), List.of(), queryParams); } // Scenario 2: ISO-8601 format parsing with extended time range @@ -4847,57 +4774,23 @@ private Stream provideFormatParsingScenarios() { @MethodSource("provideFormatParsingScenarios") void whenTimeParametersInISO8601Format_thenReturnFilteredTraces( String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var workspace = setupTestWorkspace(); Instant referenceTime = Instant.now(); Instant startTime = referenceTime.minus(Duration.ofMinutes(90)); Instant withinBoundsTime = referenceTime.minus(Duration.ofMinutes(30)); Instant outsideBoundsTime = referenceTime.plus(Duration.ofMinutes(30)); - List allTraces = new ArrayList<>(); - - // Should be included: at start of range - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(startTime)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); + List allTraces = List.of( + createTraceAtTimestamp(workspace.projectName(), startTime, "Should be included: at start of range"), + createTraceAtTimestamp(workspace.projectName(), withinBoundsTime, + "Should be included: within range"), + createTraceAtTimestamp(workspace.projectName(), referenceTime, + "Should be included: at end of range"), + createTraceAtTimestamp(workspace.projectName(), outsideBoundsTime, + "Should NOT be included: outside range")); - // Should be included: within range - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(withinBoundsTime)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - // Should be included: at end of range - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(referenceTime)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - // Should NOT be included: outside range - allTraces.add(createTrace().toBuilder() - .projectName(projectName) - .id(generateUUIDForTimestamp(outsideBoundsTime)) - .spanCount(0) - .llmSpanCount(0) - .guardrailsValidations(null) - .build()); - - traceResourceClient.batchCreateTraces(allTraces, apiKey, workspaceName); + traceResourceClient.batchCreateTraces(allTraces, workspace.apiKey(), workspace.workspaceName()); // Filter to get first 3 traces (within bounds) List expectedTraces = allTraces.subList(0, 3); @@ -4907,20 +4800,15 @@ void whenTimeParametersInISO8601Format_thenReturnFilteredTraces( "from_time", startTime.toString(), "to_time", referenceTime.toString()); - // Clear projectName from traces since API returns projectName=null - allTraces = allTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - expectedTraces = expectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - unexpectedTraces = unexpectedTraces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); - - // Sort expectedTraces by ID descending to match API response order - expectedTraces = expectedTraces.stream() - .sorted(Comparator.comparing(Trace::id).reversed()) - .toList(); + allTraces = normalizeTraces(allTraces); + expectedTraces = normalizeTraces(expectedTraces); + unexpectedTraces = normalizeTraces(unexpectedTraces); + expectedTraces = sortByIdDescending(expectedTraces); var values = testAssertion.transformTestParams(allTraces, expectedTraces, unexpectedTraces); - testAssertion.assertTest(projectName, null, apiKey, workspaceName, values.expected(), - values.unexpected(), values.all(), List.of(), queryParams); + testAssertion.assertTest(workspace.projectName(), null, workspace.apiKey(), workspace.workspaceName(), + values.expected(), values.unexpected(), values.all(), List.of(), queryParams); } // Scenario 3: Incomplete time parameters should be rejected @@ -4935,23 +4823,17 @@ private Stream provideInvalidParameterScenarios() { @MethodSource("provideInvalidParameterScenarios") void whenOnlyFromTimeProvided_thenReturnBadRequest( String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - + var workspace = setupTestWorkspace(); Instant now = Instant.now(); WebTarget target = client.target(URL_TEMPLATE.formatted(baseURI)) - .queryParam("project_name", projectName) + .queryParam("project_name", workspace.projectName()) .queryParam("from_time", now.toString()); var actualResponse = target .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) + .header(HttpHeaders.AUTHORIZATION, workspace.apiKey()) + .header(WORKSPACE_HEADER, workspace.workspaceName()) .get(); assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); @@ -4962,35 +4844,62 @@ void whenOnlyFromTimeProvided_thenReturnBadRequest( @MethodSource("provideInvalidParameterScenarios") void whenFromTimeAfterToTime_thenReturnBadRequest( String endpoint, TracePageTestAssertion testAssertion) { - var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); - var workspaceId = UUID.randomUUID().toString(); - var apiKey = UUID.randomUUID().toString(); - var projectName = RandomStringUtils.secure().nextAlphanumeric(20); - - mockTargetWorkspace(apiKey, workspaceName, workspaceId); - + var workspace = setupTestWorkspace(); Instant now = Instant.now(); Instant earlier = now.minus(Duration.ofMinutes(10)); // from_time (now) is after to_time (earlier) - should fail WebTarget target = client.target(URL_TEMPLATE.formatted(baseURI)) - .queryParam("project_name", projectName) + .queryParam("project_name", workspace.projectName()) .queryParam("from_time", now.toString()) .queryParam("to_time", earlier.toString()); var actualResponse = target .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) + .header(HttpHeaders.AUTHORIZATION, workspace.apiKey()) + .header(WORKSPACE_HEADER, workspace.workspaceName()) .get(); assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_BAD_REQUEST); } + // Helper methods to reduce duplication + private TestWorkspace setupTestWorkspace() { + var workspaceName = RandomStringUtils.secure().nextAlphanumeric(10); + var workspaceId = UUID.randomUUID().toString(); + var apiKey = UUID.randomUUID().toString(); + var projectName = RandomStringUtils.secure().nextAlphanumeric(20); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + return new TestWorkspace(workspaceName, workspaceId, apiKey, projectName); + } + + private Trace createTraceAtTimestamp(String projectName, Instant timestamp, String comment) { + return createTrace().toBuilder() + .projectName(projectName) + .id(generateUUIDForTimestamp(timestamp)) + .spanCount(0) + .llmSpanCount(0) + .guardrailsValidations(null) + .build(); + } + + private List normalizeTraces(List traces) { + return traces.stream().map(t -> t.toBuilder().projectName(null).build()).toList(); + } + + private List sortByIdDescending(List traces) { + return traces.stream() + .sorted(Comparator.comparing(Trace::id).reversed()) + .toList(); + } + private UUID generateUUIDForTimestamp(Instant timestamp) { byte[] zeroBytes = new byte[8]; return OpenTelemetryMapper.convertOtelIdToUUIDv7(zeroBytes, timestamp.toEpochMilli()); } + } } From 4a0842ab4fd553eac8c712dbccfa8c5cd0792745 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Tue, 4 Nov 2025 11:37:36 +0100 Subject: [PATCH 7/9] Revision 2: Address PR review comments - LOW priority fixes - #7: Add INSTANCE singleton pattern for InstantConverter - #8: Use StringUtils.isEmpty() for null-safe empty check - Note: #2 and #3 already addressed in previous commit --- .../comet/opik/api/InstantToUUIDMapper.java | 37 ++++++++++++++++--- .../api/resources/v1/priv/TracesResource.java | 27 ++++++-------- .../java/com/comet/opik/domain/TraceDAO.java | 18 ++++----- .../web/InstantParamConverter.java | 10 +++-- .../com/comet/opik/utils/ValidationUtils.java | 2 +- 5 files changed, 60 insertions(+), 34 deletions(-) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java index 279df93845b..6f91d56264c 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java @@ -9,7 +9,27 @@ /** * Mapper for converting Instant time boundaries to UUIDv7 bounds for efficient BETWEEN queries. + * * UUIDv7 encodes the timestamp in the first 48 bits, allowing for lexicographic sorting by time. + * This mapper creates predictable UUID boundaries (with zero random bits) to ensure correct BETWEEN + * query semantics for time-based filtering. + * + * Why not use IdGenerator.getTimeOrderedEpoch()? + * ================================================ + * While IdGenerator.getTimeOrderedEpoch() also creates UUIDs from timestamps using the reliable + * java-uuid-generator library, it uses random bits for the lower 80 bits of the UUID. For time-range + * queries, we need: + * - Lower bound: UUID with ALL ZEROS for random bits (lexicographically smallest UUID at this timestamp) + * - Upper bound: UUID at (timestamp+1ms) with ALL ZEROS (first UUID AFTER the end time) + * + * This ensures that BETWEEN queries correctly include all traces within the specified time range. + * IdGenerator's random bits would make the bounds non-deterministic, breaking BETWEEN semantics. + * + * Implementation Note: + * We use OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], ...) which: + * 1. Takes a zero-filled byte array (produces all-zero random bits via SHA-256 hash) + * 2. Encodes the timestamp in bytes 0-5 + * 3. Produces deterministic, predictable UUIDs suitable for range queries */ @UtilityClass @Slf4j @@ -17,10 +37,10 @@ public class InstantToUUIDMapper { /** * Generates a UUIDv7 lower bound from a timestamp for BETWEEN queries. - * Uses 0x00 for random bits to get the lexicographically smallest UUID with this timestamp. + * Creates the lexicographically smallest UUID with this timestamp (all zeros for random bits). * * @param timestamp the instant in time - * @return the lower bound UUIDv7 + * @return the lower bound UUIDv7 (min UUID at this timestamp) */ public static UUID toLowerBound(Instant timestamp) { if (timestamp == null) { @@ -32,11 +52,18 @@ public static UUID toLowerBound(Instant timestamp) { /** * Generates a UUIDv7 upper bound from a timestamp for BETWEEN queries. - * Uses zero bytes for the NEXT millisecond to get the first UUID after the end time. - * This ensures BETWEEN includes all UUIDs created within the end timestamp. + * Creates the first UUID AFTER the specified time (by adding 1ms and using zero random bits). + * + * This ensures BETWEEN includes all UUIDs created within the end timestamp millisecond. + * For example, if querying traces between 10:00:00.000 and 10:00:01.000: + * - toLowerBound(10:00:00.000) gives the min UUID at 10:00:00 + * - toUpperBound(10:00:01.000) gives the min UUID at 10:00:02 (next millisecond) + * - BETWEEN x AND y includes all UUIDs from 10:00:00 up to but NOT including 10:00:02 + * - Which correctly includes all UUIDs from 10:00:00 through 10:00:01.999 + * * @param timestamp the instant in time - * @return the upper bound UUIDv7 (UUID for next millisecond with zero random bits) + * @return the upper bound UUIDv7 (min UUID at timestamp+1ms) */ public static UUID toUpperBound(Instant timestamp) { if (timestamp == null) { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java index 8f403b6ee52..ee679ac7c91 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java @@ -129,8 +129,8 @@ public Response getTracesByProject( @QueryParam("strip_attachments") @DefaultValue("false") @Schema(description = "If true, returns attachment references like [file.png]; if false, downloads and reinjects stripped attachments") boolean stripAttachments, @QueryParam("sorting") String sorting, @QueryParam("exclude") String exclude, - @QueryParam("from_time") Instant startTime, - @QueryParam("to_time") Instant endTime) { + @QueryParam("from_time") @Schema(description = "Filter traces created from this time (ISO-8601 format). Must be provided together with 'to_time'.") Instant startTime, + @QueryParam("to_time") @Schema(description = "Filter traces created up to this time (ISO-8601 format). Must be provided together with 'from_time' and must be after 'from_time'.") Instant endTime) { validateProjectNameAndProjectId(projectName, projectId); validateTimeRangeParameters(startTime, endTime); @@ -174,7 +174,8 @@ public Response getTracesByProject( .contextWrite(ctx -> setRequestContext(ctx, requestContext)) .block(); - log.info("Found traces by '{}', count '{}' on workspaceId '{}'", searchCriteria, tracePage.size(), workspaceId); + log.info("Found traces by '{}', count '{}' on workspaceId '{}'", searchCriteria, tracePage.size(), + workspaceId); return Response.ok(tracePage).build(); } @@ -362,26 +363,20 @@ public Response deleteTraces( public Response getStats(@QueryParam("project_id") UUID projectId, @QueryParam("project_name") String projectName, @QueryParam("filters") String filters, - @QueryParam("from_time") Instant startTime, - @QueryParam("to_time") Instant endTime) { + @QueryParam("from_time") @Schema(description = "Filter traces created from this time (ISO-8601 format). Must be provided together with 'to_time'.") Instant startTime, + @QueryParam("to_time") @Schema(description = "Filter traces created up to this time (ISO-8601 format). Must be provided together with 'from_time' and must be after 'from_time'.") Instant endTime) { validateProjectNameAndProjectId(projectName, projectId); validateTimeRangeParameters(startTime, endTime); var traceFilters = filtersFactory.newFilters(filters, TraceFilter.LIST_TYPE_REFERENCE); - var searchCriteriaBuilder = TraceSearchCriteria.builder() + var searchCriteria = TraceSearchCriteria.builder() .projectName(projectName) .projectId(projectId) - .filters(traceFilters); - - // Apply UUID-based time filtering if both from_time and to_time are provided - if (startTime != null && endTime != null) { - searchCriteriaBuilder - .uuidFromTime(InstantToUUIDMapper.toLowerBound(startTime)) - .uuidToTime(InstantToUUIDMapper.toUpperBound(endTime)); - } - - var searchCriteria = searchCriteriaBuilder.build(); + .filters(traceFilters) + .uuidFromTime(InstantToUUIDMapper.toLowerBound(startTime)) + .uuidToTime(InstantToUUIDMapper.toUpperBound(endTime)) + .build(); String workspaceId = requestContext.get().getWorkspaceId(); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java index 194280873bb..bfdad342650 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java @@ -781,7 +781,7 @@ AND notEquals(start_time, toDateTime64('1970-01-01 00:00:00.000', 9)), WHERE workspace_id = :workspace_id AND project_id = :project_id - AND id BETWEEN :uuid_from_time AND :uuid_to_time + AND id BETWEEN :uuid_from_time AND :uuid_to_time AND id \\< :last_received_id AND AND @@ -996,7 +996,7 @@ AND notEquals(start_time, toDateTime64('1970-01-01 00:00:00.000', 9)), WHERE project_id = :project_id AND workspace_id = :workspace_id - AND id BETWEEN :uuid_from_time AND :uuid_to_time + AND id BETWEEN :uuid_from_time AND :uuid_to_time AND AND @@ -1414,7 +1414,7 @@ AND notEquals(start_time, toDateTime64('1970-01-01 00:00:00.000', 9)), WHERE workspace_id = :workspace_id AND project_id IN :project_ids - AND id BETWEEN :uuid_from_time AND :uuid_to_time + AND id BETWEEN :uuid_from_time AND :uuid_to_time AND AND @@ -2896,12 +2896,12 @@ private ST newFindTemplate(String query, TraceSearchCriteria traceSearchCriteria Optional.ofNullable(traceSearchCriteria.lastReceivedId()) .ifPresent(lastReceivedTraceId -> template.add("last_received_id", lastReceivedTraceId)); - // Bind UUID BETWEEN bounds for time-based filtering - if (traceSearchCriteria.uuidFromTime() != null && traceSearchCriteria.uuidToTime() != null) { - template.add("uuid_between", true); - template.add("uuid_from_time", traceSearchCriteria.uuidFromTime()); - template.add("uuid_to_time", traceSearchCriteria.uuidToTime()); - } + // Add UUID bounds for time-based filtering (presence of uuid_from_time triggers the conditional) + Optional.ofNullable(traceSearchCriteria.uuidFromTime()) + .ifPresent(uuid_from_time -> { + template.add("uuid_from_time", uuid_from_time); + template.add("uuid_to_time", traceSearchCriteria.uuidToTime()); + }); return template; } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java index 4c9c9fa3e1f..e68542f7f26 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/web/InstantParamConverter.java @@ -5,11 +5,13 @@ import jakarta.ws.rs.ext.ParamConverterProvider; import jakarta.ws.rs.ext.Provider; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; import java.lang.annotation.Annotation; import java.lang.reflect.Type; import java.time.Instant; import java.time.format.DateTimeParseException; +import java.util.Objects; /** * JAX-RS ParamConverter for automatic Instant conversion from query parameters. @@ -19,20 +21,22 @@ @Slf4j public class InstantParamConverter implements ParamConverterProvider { + private static final InstantConverter INSTANCE = new InstantConverter(); + @Override public ParamConverter getConverter(Class rawType, Type genericType, Annotation[] annotations) { if (rawType != Instant.class) { return null; } - return (ParamConverter) new InstantConverter(); + return (ParamConverter) INSTANCE; } private static class InstantConverter implements ParamConverter { @Override public Instant fromString(String value) { - if (value == null || value.isEmpty()) { + if (StringUtils.isEmpty(value)) { return null; } @@ -56,7 +60,7 @@ public Instant fromString(String value) { @Override public String toString(Instant value) { - return value != null ? value.toString() : null; + return Objects.toString(value, null); } } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java index eef70217c77..b97572ffb14 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/utils/ValidationUtils.java @@ -60,7 +60,7 @@ public static void validateTimeRangeParameters(Instant startTime, Instant endTim "Both 'from_time' and 'to_time' parameters must be provided together, or both must be omitted"); } - if (startTimePresent && endTimePresent && startTime.isAfter(endTime)) { + if (startTimePresent && endTimePresent && !startTime.isBefore(endTime)) { throw new BadRequestException( "Parameter 'from_time' must be before 'to_time'"); } From 4a4b3500dcb34813cbaed832c4f5f667b02e4419 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Tue, 4 Nov 2025 12:52:19 +0100 Subject: [PATCH 8/9] Revision 3: Use IdGenerator.getTimeOrderedEpoch() for UUID bounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Simplified InstantToUUIDMapper to use IdGenerator.getTimeOrderedEpoch() instead of convertOtelIdToUUIDv7 - Per UUIDv7 RFC, sub-millisecond 12 bits are optional with millisecond granularity - Start/end interval semantics with ±1ms ensures correct BETWEEN query results - This approach has been battle-tested for months without issues per reviewer recommendation - Converted InstantToUUIDMapper to @Singleton service for proper DI integration - Updated TracesResource to inject InstantToUUIDMapper dependency - Updated tests to properly mock IdGenerator dependency --- .../comet/opik/api/InstantToUUIDMapper.java | 55 +++-- .../api/resources/v1/priv/TracesResource.java | 9 +- .../opik/api/InstantToUUIDMapperTest.java | 202 ++++-------------- 3 files changed, 72 insertions(+), 194 deletions(-) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java index 6f91d56264c..05966da6f83 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/InstantToUUIDMapper.java @@ -1,7 +1,9 @@ package com.comet.opik.api; -import com.comet.opik.domain.OpenTelemetryMapper; -import lombok.experimental.UtilityClass; +import com.comet.opik.domain.IdGenerator; +import com.google.inject.Singleton; +import jakarta.inject.Inject; +import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import java.time.Instant; @@ -11,68 +13,59 @@ * Mapper for converting Instant time boundaries to UUIDv7 bounds for efficient BETWEEN queries. * * UUIDv7 encodes the timestamp in the first 48 bits, allowing for lexicographic sorting by time. - * This mapper creates predictable UUID boundaries (with zero random bits) to ensure correct BETWEEN - * query semantics for time-based filtering. - * - * Why not use IdGenerator.getTimeOrderedEpoch()? - * ================================================ - * While IdGenerator.getTimeOrderedEpoch() also creates UUIDs from timestamps using the reliable - * java-uuid-generator library, it uses random bits for the lower 80 bits of the UUID. For time-range - * queries, we need: - * - Lower bound: UUID with ALL ZEROS for random bits (lexicographically smallest UUID at this timestamp) - * - Upper bound: UUID at (timestamp+1ms) with ALL ZEROS (first UUID AFTER the end time) - * - * This ensures that BETWEEN queries correctly include all traces within the specified time range. - * IdGenerator's random bits would make the bounds non-deterministic, breaking BETWEEN semantics. + * This mapper creates UUID boundaries to ensure correct BETWEEN query semantics for time-based filtering. * * Implementation Note: - * We use OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], ...) which: - * 1. Takes a zero-filled byte array (produces all-zero random bits via SHA-256 hash) - * 2. Encodes the timestamp in bytes 0-5 - * 3. Produces deterministic, predictable UUIDs suitable for range queries + * We use IdGenerator.getTimeOrderedEpoch() which reliably creates UUIDs from timestamps. + * Per UUIDv7 RFC, the sub-millisecond 12 bits are optional and depend on implementation. + * Our implementation uses monotonic values with millisecond granularity, so using the start/end + * interval semantics with ±1ms ensures correct BETWEEN query results. + * This approach has been battle-tested for months without issues. */ -@UtilityClass +@Singleton +@RequiredArgsConstructor(onConstructor_ = @Inject) @Slf4j public class InstantToUUIDMapper { + private final IdGenerator idGenerator; + /** * Generates a UUIDv7 lower bound from a timestamp for BETWEEN queries. - * Creates the lexicographically smallest UUID with this timestamp (all zeros for random bits). + * Creates the lexicographically smallest UUID with this timestamp. * * @param timestamp the instant in time * @return the lower bound UUIDv7 (min UUID at this timestamp) */ - public static UUID toLowerBound(Instant timestamp) { + public UUID toLowerBound(Instant timestamp) { if (timestamp == null) { return null; } - return OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], timestamp.toEpochMilli()); + return idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli()); } /** * Generates a UUIDv7 upper bound from a timestamp for BETWEEN queries. - * Creates the first UUID AFTER the specified time (by adding 1ms and using zero random bits). + * Creates the UUID at the next millisecond to ensure inclusive BETWEEN semantics. * * This ensures BETWEEN includes all UUIDs created within the end timestamp millisecond. * For example, if querying traces between 10:00:00.000 and 10:00:01.000: - * - toLowerBound(10:00:00.000) gives the min UUID at 10:00:00 - * - toUpperBound(10:00:01.000) gives the min UUID at 10:00:02 (next millisecond) + * - toLowerBound(10:00:00.000) gives the UUID at 10:00:00 + * - toUpperBound(10:00:01.000) gives the UUID at 10:00:02 (next millisecond) * - BETWEEN x AND y includes all UUIDs from 10:00:00 up to but NOT including 10:00:02 * - Which correctly includes all UUIDs from 10:00:00 through 10:00:01.999 - * * @param timestamp the instant in time - * @return the upper bound UUIDv7 (min UUID at timestamp+1ms) + * @return the upper bound UUIDv7 (UUID at timestamp+1ms) */ - public static UUID toUpperBound(Instant timestamp) { + public UUID toUpperBound(Instant timestamp) { if (timestamp == null) { return null; } - // Add 1ms and use zero random bytes to get the first UUID AFTER the end time + // Add 1ms to get the first UUID AFTER the end time // BETWEEN will include all UUIDs from toLowerBound to (but not including) this value long nextMillis = timestamp.toEpochMilli() + 1; - return OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], nextMillis); + return idGenerator.getTimeOrderedEpoch(nextMillis); } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java index ee679ac7c91..885b3509eb2 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java @@ -114,6 +114,7 @@ public class TracesResource { private final @NonNull Streamer streamer; private final @NonNull ProjectService projectService; private final @NonNull TraceThreadService traceThreadService; + private final @NonNull InstantToUUIDMapper instantToUUIDMapper; @GET @Operation(operationId = "getTracesByProject", summary = "Get traces by project_name or project_id", description = "Get traces by project_name or project_id", responses = { @@ -155,8 +156,8 @@ public Response getTracesByProject( .filters(traceFilters) .truncate(truncate) .stripAttachments(stripAttachments) - .uuidFromTime(InstantToUUIDMapper.toLowerBound(startTime)) - .uuidToTime(InstantToUUIDMapper.toUpperBound(endTime)) + .uuidFromTime(instantToUUIDMapper.toLowerBound(startTime)) + .uuidToTime(instantToUUIDMapper.toUpperBound(endTime)) .exclude(ParamsValidator.get(exclude, Trace.TraceField.class, "exclude")) .sortingFields(sortingFields) .build(); @@ -374,8 +375,8 @@ public Response getStats(@QueryParam("project_id") UUID projectId, .projectName(projectName) .projectId(projectId) .filters(traceFilters) - .uuidFromTime(InstantToUUIDMapper.toLowerBound(startTime)) - .uuidToTime(InstantToUUIDMapper.toUpperBound(endTime)) + .uuidFromTime(instantToUUIDMapper.toLowerBound(startTime)) + .uuidToTime(instantToUUIDMapper.toUpperBound(endTime)) .build(); String workspaceId = requestContext.get().getWorkspaceId(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java index e9387152d90..a56b0c4d059 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java @@ -1,12 +1,16 @@ package com.comet.opik.api; -import com.comet.opik.domain.OpenTelemetryMapper; +import com.comet.opik.domain.IdGenerator; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; import java.time.Instant; import java.util.UUID; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.when; /** * Unit tests for InstantToUUIDMapper to validate UUIDv7 boundary generation @@ -14,213 +18,93 @@ */ class InstantToUUIDMapperTest { - @Test - void shouldGenerateLowerBound_withZeroBytesForRandomPortion() { - // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); - - // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - - // Then - assertThat(lowerBound).isNotNull(); - assertThat(lowerBound.toString()).matches("[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[0-9a-f]{4}-[0-9a-f]{12}"); - } - - @Test - void shouldGenerateUpperBound_withAllFFBytesForRandomPortion() { - // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); - - // When - UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); - - // Then - assertThat(upperBound).isNotNull(); - assertThat(upperBound.toString()).matches("[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[0-9a-f]{4}-[0-9a-f]{12}"); - } + @Mock + private IdGenerator idGenerator; - @Test - void shouldProduceDifferentUUIDsForZeroAndFFBytes() { - // Given - Despite trying to use 0x00 and 0xFF for random bytes, - // they get SHA-256 hashed, so we just verify they produce different UUIDs - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + private InstantToUUIDMapper mapper; - // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); - - // Then - assertThat(lowerBound).isNotEqualTo(upperBound); + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + mapper = new InstantToUUIDMapper(idGenerator); } @Test - void shouldEncodeSameTimestampInBothBounds() { + void shouldGenerateLowerBound_withValidUUID() { // Given Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + UUID expectedUUID = UUID.fromString("1926efec-0000-7000-0000-000000000000"); + when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli())).thenReturn(expectedUUID); // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + UUID lowerBound = mapper.toLowerBound(timestamp); // Then - // Extract the timestamp portion (first 48 bits / first 12 hex chars) - String lowerBoundStr = lowerBound.toString().replace("-", ""); - String upperBoundStr = upperBound.toString().replace("-", ""); - - String lowerTimestampPart = lowerBoundStr.substring(0, 12); - String upperTimestampPart = upperBoundStr.substring(0, 12); - - // Upper bound uses next millisecond (+1ms), so timestamps should differ by 1 - UUID nextMillisLowerBound = InstantToUUIDMapper.toLowerBound(timestamp.plusMillis(1)); - String nextMillisTimestampPart = nextMillisLowerBound.toString().replace("-", "").substring(0, 12); - - assertThat(upperTimestampPart) - .as("Upper bound should encode the next millisecond") - .isEqualTo(nextMillisTimestampPart); + assertThat(lowerBound).isEqualTo(expectedUUID); } @Test - void shouldHaveDifferentRandomPortions() { + void shouldGenerateUpperBound_withIncrementedTime() { // Given Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + UUID expectedUUID = UUID.fromString("1926efec-0000-7000-0000-000000000001"); + long nextMillis = timestamp.toEpochMilli() + 1; + when(idGenerator.getTimeOrderedEpoch(nextMillis)).thenReturn(expectedUUID); // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); + UUID upperBound = mapper.toUpperBound(timestamp); // Then - // Lower bound uses same timestamp with zero random bytes - // Upper bound uses next millisecond with zero random bytes - // Since they have different timestamps, they should be different UUIDs - assertThat(lowerBound) - .as("Upper and lower bounds should be different UUIDs") - .isNotEqualTo(upperBound); + assertThat(upperBound).isEqualTo(expectedUUID); } @Test - void shouldMaintainChronologicalOrder() { - // Given - Instant time1 = Instant.parse("2025-01-15T10:00:00Z"); - Instant time2 = Instant.parse("2025-01-15T11:00:00Z"); - Instant time3 = Instant.parse("2025-01-15T12:00:00Z"); - - // When - UUID lower1 = InstantToUUIDMapper.toLowerBound(time1); - UUID lower2 = InstantToUUIDMapper.toLowerBound(time2); - UUID lower3 = InstantToUUIDMapper.toLowerBound(time3); - - // Then - assertThat(lower1.compareTo(lower2)).isNegative(); - assertThat(lower2.compareTo(lower3)).isNegative(); - assertThat(lower1.compareTo(lower3)).isNegative(); - } - - @Test - void shouldReturnNull_whenTimestampIsNull() { + void shouldReturnNull_whenLowerBoundTimestampIsNull() { // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(null); - UUID upperBound = InstantToUUIDMapper.toUpperBound(null); + UUID lowerBound = mapper.toLowerBound(null); // Then assertThat(lowerBound).isNull(); - assertThat(upperBound).isNull(); } @Test - void shouldProduceUUIDv7Format() { - // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); - + void shouldReturnNull_whenUpperBoundTimestampIsNull() { // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); + UUID upperBound = mapper.toUpperBound(null); // Then - // UUIDv7 has version nibble = 7 in the 13th character - String uuidStr = lowerBound.toString(); - assertThat(uuidStr.charAt(14)).isEqualTo('7'); - } - - @Test - void shouldGenerateConsistentUUIDs() { - // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); - - // When - UUID lowerBound1 = InstantToUUIDMapper.toLowerBound(timestamp); - UUID lowerBound2 = InstantToUUIDMapper.toLowerBound(timestamp); - - // Then - assertThat(lowerBound1).isEqualTo(lowerBound2); - } - - @Test - void shouldUseZeroBytesForLowerBound() { - // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); - - // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - UUID referenceUUID = OpenTelemetryMapper.convertOtelIdToUUIDv7(new byte[8], timestamp.toEpochMilli()); - - // Then - assertThat(lowerBound).isEqualTo(referenceUUID); + assertThat(upperBound).isNull(); } @Test - void shouldUseFFBytesForUpperBound() { + void shouldUseSameGeneratorForBothBounds() { // Given Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + UUID lowerUUID = UUID.fromString("1926efec-0000-7000-0000-000000000000"); + UUID upperUUID = UUID.fromString("1926efec-0000-7000-0000-000000000001"); - // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - UUID upperBound = InstantToUUIDMapper.toUpperBound(timestamp); - - // Then - assertThat(lowerBound).isNotNull(); - assertThat(upperBound).isNotNull(); - - // Upper bound should be greater than lower bound - assertThat(upperBound.toString().compareTo(lowerBound.toString())) - .as("Upper bound should be lexicographically greater than lower bound") - .isGreaterThan(0); - } - - @Test - void shouldWorkWithBoundaryInstants() { - // Given - Test with epoch start - Instant epochStart = Instant.EPOCH; + when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli())).thenReturn(lowerUUID); + when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli() + 1)).thenReturn(upperUUID); // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(epochStart); - UUID upperBound = InstantToUUIDMapper.toUpperBound(epochStart); + UUID lower = mapper.toLowerBound(timestamp); + UUID upper = mapper.toUpperBound(timestamp); // Then - assertThat(lowerBound).isNotNull(); - assertThat(upperBound).isNotNull(); - // They should be different due to different random bytes (hashed) - assertThat(lowerBound).isNotEqualTo(upperBound); + assertThat(lower).isEqualTo(lowerUUID); + assertThat(upper).isEqualTo(upperUUID); } @Test - void shouldAllUUIDsWithSameTimestampShareTimestampPortion() { + void shouldIncrementTimeByOneMillisecond_forUpperBound() { // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + Instant timestamp = Instant.parse("2025-01-15T10:30:00.500Z"); + long expectedUpperMillis = timestamp.toEpochMilli() + 1; // When - UUID lowerBound = InstantToUUIDMapper.toLowerBound(timestamp); - UUID nextMillisLowerBound = InstantToUUIDMapper.toLowerBound(timestamp.plusMillis(1)); - - // Then - // Extract the timestamp portion (first 48 bits / first 12 hex chars) - String lowerBoundStr = lowerBound.toString().replace("-", ""); - String nextMillisStr = nextMillisLowerBound.toString().replace("-", ""); - - String lowerTimestampPart = lowerBoundStr.substring(0, 12); - String nextMillisTimestampPart = nextMillisStr.substring(0, 12); + mapper.toUpperBound(timestamp); - assertThat(nextMillisTimestampPart) - .as("Next millisecond should have different timestamp portion") - .isNotEqualTo(lowerTimestampPart); + // Then - Verify the IdGenerator was called with the incremented time + org.mockito.Mockito.verify(idGenerator).getTimeOrderedEpoch(expectedUpperMillis); } } From bdd2249f030a29675eab58bc36e209d94fae2237 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Tue, 4 Nov 2025 13:30:03 +0100 Subject: [PATCH 9/9] Fix tests --- .../opik/api/InstantToUUIDMapperTest.java | 127 +++++++++++++++--- 1 file changed, 105 insertions(+), 22 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java index a56b0c4d059..0f2c3639da3 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/InstantToUUIDMapperTest.java @@ -10,11 +10,12 @@ import java.util.UUID; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; /** * Unit tests for InstantToUUIDMapper to validate UUIDv7 boundary generation - * and ensure consistent timestamp encoding. + * and ensure consistent timestamp encoding for time-based filtering. */ class InstantToUUIDMapperTest { @@ -30,32 +31,35 @@ void setUp() { } @Test - void shouldGenerateLowerBound_withValidUUID() { + void shouldGenerateLowerBound_withTimestampEncoded() { // Given Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + long expectedMillis = timestamp.toEpochMilli(); UUID expectedUUID = UUID.fromString("1926efec-0000-7000-0000-000000000000"); - when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli())).thenReturn(expectedUUID); + when(idGenerator.getTimeOrderedEpoch(expectedMillis)).thenReturn(expectedUUID); // When UUID lowerBound = mapper.toLowerBound(timestamp); // Then assertThat(lowerBound).isEqualTo(expectedUUID); + verify(idGenerator).getTimeOrderedEpoch(expectedMillis); } @Test - void shouldGenerateUpperBound_withIncrementedTime() { + void shouldGenerateUpperBound_withNextMillisecond() { // Given Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + long expectedMillis = timestamp.toEpochMilli() + 1; UUID expectedUUID = UUID.fromString("1926efec-0000-7000-0000-000000000001"); - long nextMillis = timestamp.toEpochMilli() + 1; - when(idGenerator.getTimeOrderedEpoch(nextMillis)).thenReturn(expectedUUID); + when(idGenerator.getTimeOrderedEpoch(expectedMillis)).thenReturn(expectedUUID); // When UUID upperBound = mapper.toUpperBound(timestamp); // Then assertThat(upperBound).isEqualTo(expectedUUID); + verify(idGenerator).getTimeOrderedEpoch(expectedMillis); } @Test @@ -77,34 +81,113 @@ void shouldReturnNull_whenUpperBoundTimestampIsNull() { } @Test - void shouldUseSameGeneratorForBothBounds() { - // Given - Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); - UUID lowerUUID = UUID.fromString("1926efec-0000-7000-0000-000000000000"); - UUID upperUUID = UUID.fromString("1926efec-0000-7000-0000-000000000001"); + void shouldMaintainTimestampBoundarySemantics() { + // Given - Time range query semantics: BETWEEN lower AND upper + Instant startTime = Instant.parse("2025-01-15T10:30:00.000Z"); + Instant endTime = Instant.parse("2025-01-15T10:30:01.000Z"); + + UUID startLower = UUID.fromString("1926efec-0000-7000-0000-000000000000"); + UUID endUpper = UUID.fromString("1926efed-0000-7000-0000-000000000000"); - when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli())).thenReturn(lowerUUID); - when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli() + 1)).thenReturn(upperUUID); + when(idGenerator.getTimeOrderedEpoch(startTime.toEpochMilli())).thenReturn(startLower); + when(idGenerator.getTimeOrderedEpoch(endTime.toEpochMilli() + 1)).thenReturn(endUpper); // When - UUID lower = mapper.toLowerBound(timestamp); - UUID upper = mapper.toUpperBound(timestamp); + UUID lower = mapper.toLowerBound(startTime); + UUID upper = mapper.toUpperBound(endTime); - // Then - assertThat(lower).isEqualTo(lowerUUID); - assertThat(upper).isEqualTo(upperUUID); + // Then - Verify the bounds correctly represent the time range + assertThat(lower).isEqualTo(startLower); + assertThat(upper).isEqualTo(endUpper); + + // Verify IdGenerator was called with correct millisecond values + verify(idGenerator).getTimeOrderedEpoch(startTime.toEpochMilli()); + verify(idGenerator).getTimeOrderedEpoch(endTime.toEpochMilli() + 1); } @Test - void shouldIncrementTimeByOneMillisecond_forUpperBound() { + void shouldAddOneMillisecondToUpperBoundTime() { // Given Instant timestamp = Instant.parse("2025-01-15T10:30:00.500Z"); - long expectedUpperMillis = timestamp.toEpochMilli() + 1; + long lowerMillis = timestamp.toEpochMilli(); + long upperMillis = lowerMillis + 1; // When + mapper.toLowerBound(timestamp); mapper.toUpperBound(timestamp); - // Then - Verify the IdGenerator was called with the incremented time - org.mockito.Mockito.verify(idGenerator).getTimeOrderedEpoch(expectedUpperMillis); + // Then - Verify that upper bound is exactly 1ms after lower bound + verify(idGenerator).getTimeOrderedEpoch(lowerMillis); + verify(idGenerator).getTimeOrderedEpoch(upperMillis); + assertThat(upperMillis).isEqualTo(lowerMillis + 1); + } + + @Test + void shouldHandleEpochTime() { + // Given + Instant epochTime = Instant.EPOCH; + UUID expectedUUID = UUID.fromString("00000000-0000-7000-0000-000000000000"); + when(idGenerator.getTimeOrderedEpoch(0L)).thenReturn(expectedUUID); + + // When + UUID lowerBound = mapper.toLowerBound(epochTime); + + // Then + assertThat(lowerBound).isEqualTo(expectedUUID); + verify(idGenerator).getTimeOrderedEpoch(0L); + } + + @Test + void shouldHandleLargeTimestamps() { + // Given + Instant futureTime = Instant.parse("2099-12-31T23:59:59.999Z"); + long expectedMillis = futureTime.toEpochMilli(); + UUID expectedUUID = UUID.fromString("ffffffff-ffff-7fff-ffff-ffffffffffff"); + when(idGenerator.getTimeOrderedEpoch(expectedMillis)).thenReturn(expectedUUID); + + // When + UUID lowerBound = mapper.toLowerBound(futureTime); + + // Then + assertThat(lowerBound).isEqualTo(expectedUUID); + verify(idGenerator).getTimeOrderedEpoch(expectedMillis); + } + + @Test + void shouldUseSameGeneratorForConsistentResults() { + // Given + Instant timestamp = Instant.parse("2025-01-15T10:30:00Z"); + UUID firstCallUUID = UUID.fromString("1926efec-0000-7000-0000-000000000000"); + when(idGenerator.getTimeOrderedEpoch(timestamp.toEpochMilli())).thenReturn(firstCallUUID); + + // When - Call toLowerBound twice with same timestamp + UUID result1 = mapper.toLowerBound(timestamp); + UUID result2 = mapper.toLowerBound(timestamp); + + // Then - Should produce same UUID due to IdGenerator consistency + assertThat(result1).isEqualTo(result2); + assertThat(result1).isEqualTo(firstCallUUID); + } + + @Test + void shouldProduceDifferentUUIDsForDifferentTimestamps() { + // Given + Instant time1 = Instant.parse("2025-01-15T10:30:00.000Z"); + Instant time2 = Instant.parse("2025-01-15T10:30:01.000Z"); + + UUID uuid1 = UUID.fromString("1926efec-0000-7000-0000-000000000000"); + UUID uuid2 = UUID.fromString("1926efed-0000-7000-0000-000000000000"); + + when(idGenerator.getTimeOrderedEpoch(time1.toEpochMilli())).thenReturn(uuid1); + when(idGenerator.getTimeOrderedEpoch(time2.toEpochMilli())).thenReturn(uuid2); + + // When + UUID result1 = mapper.toLowerBound(time1); + UUID result2 = mapper.toLowerBound(time2); + + // Then + assertThat(result1).isNotEqualTo(result2); + assertThat(result1).isEqualTo(uuid1); + assertThat(result2).isEqualTo(uuid2); } }