/ libs / typescript / jest.global-server-setup.ts
jest.global-server-setup.ts
 1  import { spawn } from 'child_process';
 2  import { tmpdir } from 'os';
 3  import { mkdtempSync } from 'fs';
 4  import { join } from 'path';
 5  import { TEST_PORT, TEST_TRACKING_URI } from './core/tests/helper';
 6  
 7  /**
 8   * Start MLflow Python server. This is necessary for testing Typescript SDK because
 9   * the SDK does not have a server implementation and talks to the Python server instead.
10   */
11  module.exports = async () => {
12    const tempDir = mkdtempSync(join(tmpdir(), 'mlflow-test-'));
13  
14    const mlflowRoot = join(__dirname, '../..'); // Use the local dev version
15  
16    // Only start a server if one is not already running
17    try {
18      const response = await fetch(TEST_TRACKING_URI);
19      if (response.ok) {
20        return;
21      }
22    } catch (error) {
23      // Ignore error
24    }
25  
26    // eslint-disable-next-line no-console
27    console.log(`Starting MLflow server on port ${TEST_PORT}. This may take a few seconds...
28        To speed up the test, you can manually start the server and keep it running during local development.`);
29  
30    const mlflowProcess = spawn(
31      'uv',
32      ['run', '--directory', mlflowRoot, 'mlflow', 'server', '--port', TEST_PORT.toString()],
33      {
34        cwd: tempDir,
35        stdio: 'inherit',
36        // Create a new process group so we can kill the entire group
37        detached: true,
38        env: {
39          ...process.env,
40          // Disable job execution to avoid timing issues in tests
41          MLFLOW_SERVER_ENABLE_JOB_EXECUTION: 'false',
42        },
43      },
44    );
45  
46    try {
47      await waitForServer(TEST_PORT);
48      // eslint-disable-next-line no-console
49      console.log(`MLflow server is ready on port ${TEST_PORT}`);
50    } catch (error) {
51      console.error('Failed to start MLflow server:', error);
52      throw error;
53    }
54  
55    // Set global variables for cleanup in jest.global-teardown.ts
56    const globals = globalThis as any;
57    globals.mlflowProcess = mlflowProcess;
58    globals.tempDir = tempDir;
59  };
60  
61  async function waitForServer(maxAttempts: number = 30): Promise<void> {
62    for (let i = 0; i < maxAttempts; i++) {
63      try {
64        const response = await fetch(TEST_TRACKING_URI);
65        if (response.ok) {
66          return;
67        }
68      } catch (error) {
69        // Ignore error
70      }
71      await new Promise((resolve) => setTimeout(resolve, 1000));
72    }
73    throw new Error('Failed to start MLflow server');
74  }