ci(llmobs): use plugin test handler for receiving both llm observability and apm spans (#6477)

* most tests updated

* skip problematic test

* working solution for distinct traces

* revert agent changes and do something different in useLlmObs

* review comments

Author: sabrenner

packages/dd-trace/test/llmobs/plugins/google-cloud-vertexai/index.spec.js

@@ -4,12 +4,11 @@ const { expect } = require('chai')
 const { describe, it, beforeEach, afterEach, before, after } = require('mocha')
 const sinon = require('sinon')
 
-const LLMObsSpanWriter = require('../../../../src/llmobs/writers/spans')
-const agent = require('../../../plugins/agent')
 const { withVersions } = require('../../../setup/mocha')
 const {
   expectedLLMObsLLMSpanEvent,
-  deepEqualWithMockValues
+  deepEqualWithMockValues,
+  useLlmObs
 } = require('../../util')
 const chai = require('chai')
 
@@ -82,28 +81,7 @@ describe('integrations', () => {
   }
 
   describe('vertexai', () => {
-    before(async () => {
-      sinon.stub(LLMObsSpanWriter.prototype, 'append')
-
-      // reduce errors related to too many listeners
-      process.removeAllListeners('beforeExit')
-
-      return agent.load('google-cloud-vertexai', {}, {
-        llmobs: {
-          mlApp: 'test',
-          agentlessEnabled: false
-        }
-      })
-    })
-
-    afterEach(() => {
-      LLMObsSpanWriter.prototype.append.reset()
-    })
-
-    after(() => {
-      sinon.restore()
-      return agent.close({ ritmReset: false, wipe: true })
-    })
+    const getEvents = useLlmObs({ plugin: 'google-cloud-vertexai' })
 
     withVersions('google-cloud-vertexai', '@google-cloud/vertexai', '>=1', version => {
       before(() => {
@@ -137,88 +115,76 @@ describe('integrations', () => {
         useScenario({ scenario: 'generate-content-single-response' })
 
         it('makes a successful call', async () => {
-          const checkTraces = agent.assertSomeTraces(traces => {
-            const span = traces[0][0]
-            const spanEvent = LLMObsSpanWriter.prototype.append.getCall(0).args[0]
-
-            const expected = expectedLLMObsLLMSpanEvent({
-              span,
-              spanKind: 'llm',
-              modelName: 'gemini-1.5-flash-002',
-              modelProvider: 'google',
-              name: 'GenerativeModel.generateContent',
-              inputMessages: getInputMessages('Hello, how are you?'),
-              outputMessages: [
-                {
-                  role: 'model',
-                  content: 'Hello! How can I assist you today?'
-                }
-              ],
-              metadata: {
-                temperature: 1,
-                max_output_tokens: 50
-              },
-              tokenMetrics: { input_tokens: 35, output_tokens: 2, total_tokens: 37 },
-              tags: { ml_app: 'test', language: 'javascript', integration: 'vertexai' }
-            })
-
-            expect(spanEvent).to.deepEqualWithMockValues(expected)
-          })
-
           await model.generateContent({
             contents: [{ role: 'user', parts: [{ text: 'Hello, how are you?' }] }]
           })
 
-          await checkTraces
+          const { apmSpans, llmobsSpans } = await getEvents()
+          const expected = expectedLLMObsLLMSpanEvent({
+            span: apmSpans[0],
+            spanKind: 'llm',
+            modelName: 'gemini-1.5-flash-002',
+            modelProvider: 'google',
+            name: 'GenerativeModel.generateContent',
+            inputMessages: getInputMessages('Hello, how are you?'),
+            outputMessages: [
+              {
+                role: 'model',
+                content: 'Hello! How can I assist you today?'
+              }
+            ],
+            metadata: {
+              temperature: 1,
+              max_output_tokens: 50
+            },
+            tokenMetrics: { input_tokens: 35, output_tokens: 2, total_tokens: 37 },
+            tags: { ml_app: 'test', language: 'javascript', integration: 'vertexai' }
+          })
+
+          expect(llmobsSpans[0]).to.deepEqualWithMockValues(expected)
         })
       })
 
       describe('tool calls', () => {
         useScenario({ scenario: 'generate-content-single-response-with-tools' })
 
         it('makes a successful call', async () => {
-          const checkTraces = agent.assertSomeTraces(traces => {
-            const span = traces[0][0]
-            const spanEvent = LLMObsSpanWriter.prototype.append.getCall(0).args[0]
-
-            const expected = expectedLLMObsLLMSpanEvent({
-              span,
-              spanKind: 'llm',
-              modelName: 'gemini-1.5-flash-002',
-              modelProvider: 'google',
-              name: 'GenerativeModel.generateContent',
-              inputMessages: getInputMessages('what is 2 + 2?'),
-              outputMessages: [
-                {
-                  role: 'model',
-                  content: '',
-                  tool_calls: [
-                    {
-                      name: 'add',
-                      arguments: {
-                        a: 2,
-                        b: 2
-                      }
-                    }
-                  ]
-                }
-              ],
-              metadata: {
-                temperature: 1,
-                max_output_tokens: 50
-              },
-              tokenMetrics: { input_tokens: 20, output_tokens: 3, total_tokens: 23 },
-              tags: { ml_app: 'test', language: 'javascript', integration: 'vertexai' }
-            })
-
-            expect(spanEvent).to.deepEqualWithMockValues(expected)
-          })
-
           await model.generateContent({
             contents: [{ role: 'user', parts: [{ text: 'what is 2 + 2?' }] }]
           })
 
-          await checkTraces
+          const { apmSpans, llmobsSpans } = await getEvents()
+          const expected = expectedLLMObsLLMSpanEvent({
+            span: apmSpans[0],
+            spanKind: 'llm',
+            modelName: 'gemini-1.5-flash-002',
+            modelProvider: 'google',
+            name: 'GenerativeModel.generateContent',
+            inputMessages: getInputMessages('what is 2 + 2?'),
+            outputMessages: [
+              {
+                role: 'model',
+                content: '',
+                tool_calls: [
+                  {
+                    name: 'add',
+                    arguments: {
+                      a: 2,
+                      b: 2
+                    }
+                  }
+                ]
+              }
+            ],
+            metadata: {
+              temperature: 1,
+              max_output_tokens: 50
+            },
+            tokenMetrics: { input_tokens: 20, output_tokens: 3, total_tokens: 23 },
+            tags: { ml_app: 'test', language: 'javascript', integration: 'vertexai' }
+          })
+
+          expect(llmobsSpans[0]).to.deepEqualWithMockValues(expected)
         })
       })
 
@@ -227,44 +193,6 @@ describe('integrations', () => {
           useScenario({ scenario: 'generate-content-single-response' })
 
           it('makes a successful call', async () => {
-            const checkTraces = agent.assertSomeTraces(traces => {
-              const span = traces[0][0]
-              const spanEvent = LLMObsSpanWriter.prototype.append.getCall(0).args[0]
-
-              const inputMessages = []
-
-              if (model.systemInstruction) {
-                inputMessages.push({ role: 'system', content: 'Please provide an answer' })
-              }
-
-              inputMessages.push({ role: 'user', content: 'Foobar?' })
-              inputMessages.push({ role: 'model', content: 'Foobar!' })
-              inputMessages.push({ content: 'Hello, how are you?' })
-
-              const expected = expectedLLMObsLLMSpanEvent({
-                span,
-                spanKind: 'llm',
-                modelName: 'gemini-1.5-flash-002',
-                modelProvider: 'google',
-                name: 'ChatSession.sendMessage',
-                inputMessages,
-                outputMessages: [
-                  {
-                    role: 'model',
-                    content: 'Hello! How can I assist you today?'
-                  }
-                ],
-                metadata: {
-                  temperature: 1,
-                  max_output_tokens: 50
-                },
-                tokenMetrics: { input_tokens: 35, output_tokens: 2, total_tokens: 37 },
-                tags: { ml_app: 'test', language: 'javascript', integration: 'vertexai' }
-              })
-
-              expect(spanEvent).to.deepEqualWithMockValues(expected)
-            })
-
             const chat = model.startChat({
               history: [
                 { role: 'user', parts: [{ text: 'Foobar?' }] },
@@ -274,7 +202,40 @@ describe('integrations', () => {
 
             await chat.sendMessage([{ text: 'Hello, how are you?' }])
 
-            await checkTraces
+            const { apmSpans, llmobsSpans } = await getEvents()
+
+            const inputMessages = []
+
+            if (model.systemInstruction) {
+              inputMessages.push({ role: 'system', content: 'Please provide an answer' })
+            }
+
+            inputMessages.push({ role: 'user', content: 'Foobar?' })
+            inputMessages.push({ role: 'model', content: 'Foobar!' })
+            inputMessages.push({ content: 'Hello, how are you?' })
+
+            const expected = expectedLLMObsLLMSpanEvent({
+              span: apmSpans[0],
+              spanKind: 'llm',
+              modelName: 'gemini-1.5-flash-002',
+              modelProvider: 'google',
+              name: 'ChatSession.sendMessage',
+              inputMessages,
+              outputMessages: [
+                {
+                  role: 'model',
+                  content: 'Hello! How can I assist you today?'
+                }
+              ],
+              metadata: {
+                temperature: 1,
+                max_output_tokens: 50
+              },
+              tokenMetrics: { input_tokens: 35, output_tokens: 2, total_tokens: 37 },
+              tags: { ml_app: 'test', language: 'javascript', integration: 'vertexai' }
+            })
+
+            expect(llmobsSpans[0]).to.deepEqualWithMockValues(expected)
           })
         })
       })