lobe-chat

1
// @vitest-environment node
2
import { Langfuse } from 'langfuse';
3
import { CreateLangfuseTraceBody } from 'langfuse-core';
4
import { describe, expect, it, vi } from 'vitest';
5

6
import * as server from '@/config/langfuse';
7

8
import { TraceClient } from './index';
9

10
describe('TraceClient', () => {
11
  it('should not initialize Langfuse client when ENABLE_LANGFUSE is false', () => {
12
    vi.spyOn(server, 'getLangfuseConfig').mockReturnValue({
13
      ENABLE_LANGFUSE: false,
14
    } as any);
15
    const client = new TraceClient();
16
    expect(client['_client']).toBeUndefined();
17
  });
18

19
  it('should throw error if LANGFUSE keys are missing', () => {
20
    vi.spyOn(server, 'getLangfuseConfig').mockReturnValue({
21
      ENABLE_LANGFUSE: true,
22
    } as any);
23
    expect(() => new TraceClient()).toThrow('NO_LANGFUSE_KEY_ERROR');
24
  });
25

26
  it('should call trace method of Langfuse client', () => {
27
    const mockTrace = vi.fn();
28

29
    vi.spyOn(Langfuse.prototype, 'trace').mockImplementation(mockTrace);
30

31
    vi.spyOn(server, 'getLangfuseConfig').mockReturnValue({
32
      ENABLE_LANGFUSE: true,
33
      LANGFUSE_PUBLIC_KEY: 'public-key',
34
      LANGFUSE_SECRET_KEY: 'secret-key',
35
      LANGFUSE_HOST: 'host',
36
    } as any);
37

38
    const client = new TraceClient();
39
    const traceParam: CreateLangfuseTraceBody = { id: 'abc' };
40
    client.createTrace(traceParam);
41

42
    expect(mockTrace).toHaveBeenCalledWith(traceParam);
43
  });
44

45
  it('should call shutdownAsync method of Langfuse client', async () => {
46
    const mockShutdownAsync = vi.fn();
47

48
    vi.spyOn(Langfuse.prototype, 'shutdownAsync').mockImplementation(mockShutdownAsync);
49
    vi.spyOn(server, 'getLangfuseConfig').mockReturnValue({
50
      ENABLE_LANGFUSE: true,
51
      LANGFUSE_PUBLIC_KEY: 'public-key',
52
      LANGFUSE_SECRET_KEY: 'secret-key',
53
      LANGFUSE_HOST: 'host',
54
    } as any);
55

56
    const client = new TraceClient();
57
    await client.shutdownAsync();
58

59
    expect(mockShutdownAsync).toHaveBeenCalled();
60
  });
61
});
62