@Test
    public void testWorkflow() {
        when(vertx.executeBlocking(any(Callable.class), eq(false))).thenAnswer(invocation -> {
            Callable<?> callable = invocation.getArgument(0);
            return Future.succeededFuture(callable.call());
        });

        final String traceId = "trace-id";
        ProxyContext chatBackend = mock(ProxyContext.class);
        when(chatBackend.getSpanId()).thenReturn("chat");
        when(chatBackend.getTraceId()).thenReturn(traceId);
        when(chatBackend.getApiKeyData()).thenReturn(new ApiKeyData());

        // chat calls app -> core starts span
        tracker.startSpan(chatBackend);


        ProxyContext app = mock(ProxyContext.class);
        when(app.getSpanId()).thenReturn("app");
        when(app.getTraceId()).thenReturn(traceId);
        when(app.getParentSpanId()).thenReturn("chat");
        ApiKeyData apiKeyData = new ApiKeyData();
        apiKeyData.setPerRequestKey("key");
        when(app.getApiKeyData()).thenReturn(apiKeyData);

        // app calls model -> core starts span
        tracker.startSpan(app);

        TokenUsage modelTokenUsage = new TokenUsage();
        modelTokenUsage.setTotalTokens(100);
        modelTokenUsage.setCompletionTokens(80);
        modelTokenUsage.setPromptTokens(20);
        modelTokenUsage.setCost(new BigDecimal("10.0"));
        modelTokenUsage.setAggCost(new BigDecimal("10.0"));

        // core receives response from model
        when(app.getTokenUsage()).thenReturn(modelTokenUsage);

        tracker.updateModelStats(app);

        // core ends span for request to model
        tracker.endSpan(app);

        // core receives response from app
        Future<TokenUsage> appStatsFuture = tracker.getTokenStats(chatBackend);
        assertNotNull(appStatsFuture);
        TokenUsage tokenUsage = appStatsFuture.result();
        assertEquals(100, tokenUsage.getTotalTokens());
        assertEquals(80, tokenUsage.getCompletionTokens());
        assertEquals(20, tokenUsage.getPromptTokens());
        assertEquals(new BigDecimal("10.0"), tokenUsage.getAggCost());
        assertNull(tokenUsage.getCost());

        // core ends span for request to app
        tracker.endSpan(chatBackend);
        assertNull(tracker.getTokenStats(chatBackend).result());
    }
Tests the flow: chat back-end -> core -> app -> core -> model
server/src/test/java/com/epam/aidial/core/server/token/TokenStatsTrackerTest.java