/*
 * Copyright 2025 Adobe. All rights reserved.
 * This file is licensed to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under
 * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
 * OF ANY KIND, either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */
/** biome-ignore-all lint/style/noMagicNumbers: Test file. */

import {
  metrics,
  ROOT_CONTEXT,
  SpanStatusCode,
  trace,
} from "@opentelemetry/api";
import { NodeSDK } from "@opentelemetry/sdk-node";
import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";

import type { Context, Meter, SpanContext, Tracer } from "@opentelemetry/api";
import type { Mock, MockInstance } from "vitest";
import type { InstrumentationContext, TelemetryConfig } from "#types";

describe("core/instrumentation", () => {
  let instrumentation: typeof import("#core/instrumentation");
  let runtimeHelpers: typeof import("#helpers/runtime");

  let tracer: Tracer;
  let meter: Meter;

  beforeEach(async () => {
    vi.clearAllMocks();
    vi.resetModules();

    tracer = trace.getTracer("test-tracer");
    meter = metrics.getMeter("test-meter");
    vi.stubGlobal("__OTEL_TELEMETRY_API__", {
      meter,
      tracer,
    });

    vi.doMock("#helpers/runtime", () => ({
      getRuntimeActionMetadata: vi.fn(() => ({
        actionName: "test-action",
        actionVersion: "1.0.0",
        isDevelopment: false,
        namespace: "test-namespace",
        packageName: "test-package",
      })),
      getRuntimeInvocationAttributes: vi.fn(() => ({
        "action.activation_id": "test-activation",
        "action.region": "test-region",
      })),
      isDevelopment: vi.fn(() => false),

      isTelemetryEnabled: vi.fn(() => true),
    }));

    instrumentation = await import("#core/instrumentation");
    runtimeHelpers = await import("#helpers/runtime");
  });

  afterEach(() => {
    vi.unstubAllGlobals();
  });

  describe("getInstrumentationHelpers", () => {
    beforeEach(() => {
      vi.stubGlobal("__OTEL_SDK__", new NodeSDK());
    });

    test("should throw different errors when telemetry is not enabled or called outside instrumented function", () => {
      // This should throw because telemetry is not enabled
      vi.mocked(runtimeHelpers.isTelemetryEnabled).mockReturnValue(false);
      expect(() => instrumentation.getInstrumentationHelpers()).toThrow();

      let telemetryNotEnabledError: Error | null = null;
      try {
        instrumentation.getInstrumentationHelpers();
      } catch (error) {
        telemetryNotEnabledError = error as Error;
      }

      // This should also throw because we're outside an instrumented function
      vi.mocked(runtimeHelpers.isTelemetryEnabled).mockReturnValue(true);
      expect(() => instrumentation.getInstrumentationHelpers()).toThrow();

      let outsideInstrumentedFunctionError: Error | null = null;
      try {
        instrumentation.getInstrumentationHelpers();
      } catch (error) {
        outsideInstrumentedFunctionError = error as Error;
      }

      expect(telemetryNotEnabledError).not.toBe(
        outsideInstrumentedFunctionError,
      );

      expect(telemetryNotEnabledError?.message).not.toEqual(
        outsideInstrumentedFunctionError?.message,
      );
    });

    test("should return helpers when called within instrumented function", () => {
      let capturedHelpers: InstrumentationContext | null = null;
      const testFn = instrumentation.instrument(function testFunction() {
        capturedHelpers = instrumentation.getInstrumentationHelpers();
        return "success";
      });

      const result = testFn();
      expect(result).toBe("success");

      if (!capturedHelpers) {
        expect.fail("capturedHelpers is null");
      }

      const helpers = capturedHelpers as InstrumentationContext;
      expect(helpers).toBeDefined();
      expect(helpers.currentSpan).toBeDefined();
      expect(helpers.logger).toBeDefined();
      expect(helpers.tracer).toBeDefined();
      expect(helpers.meter).toBeDefined();
      expect(helpers.contextCarrier).toBeDefined();
    });
  });

  describe("instrument", () => {
    beforeEach(() => {
      vi.stubGlobal("__OTEL_SDK__", new NodeSDK());
    });

    test("should wrap a function and return the same result", () => {
      const originalFn = vi.fn((a: number, b: number) => a + b);
      const instrumentedFn = instrumentation.instrument(originalFn);
      const result = instrumentedFn(2, 3);

      expect(result).toBe(5);
      expect(originalFn).toHaveBeenCalledWith(2, 3);
    });

    test("should attach current invocation attributes to each span", () => {
      const spy = vi.spyOn(tracer, "startActiveSpan");
      const instrumentedFn = instrumentation.instrument(
        vi.fn(function testFn() {
          return "result";
        }),
        {
          spanConfig: {
            attributes: {
              custom: "value",
            },
          },
        },
      );

      vi.mocked(runtimeHelpers.getRuntimeInvocationAttributes)
        .mockReturnValueOnce({
          "action.activation_id": "first-activation",
          "action.region": "first-region",
        })
        .mockReturnValueOnce({
          "action.activation_id": "second-activation",
          "action.region": "second-region",
        });

      instrumentedFn();
      instrumentedFn();

      expect(spy.mock.calls[0][1]).toMatchObject({
        attributes: {
          "action.activation_id": "first-activation",
          "action.region": "first-region",
          custom: "value",
        },
      });
      expect(spy.mock.calls[1][1]).toMatchObject({
        attributes: {
          "action.activation_id": "second-activation",
          "action.region": "second-region",
          custom: "value",
        },
      });
    });

    test("should use function name as span name by default", () => {
      const spy = vi.spyOn(tracer, "startActiveSpan");
      const namedFunction = function myFunction() {
        return "result";
      };

      const instrumented = instrumentation.instrument(namedFunction);
      instrumented();

      expect(spy.mock.calls[0][0]).toBe("myFunction");
    });

    test("should use provided span name", () => {
      const spy = vi.spyOn(tracer, "startActiveSpan");
      const fn = vi.fn(function testFn() {
        return "result";
      });

      const instrumented1 = instrumentation.instrument(fn, {
        spanConfig: { spanName: "custom-span" },
      });

      expect(instrumented1()).toBe("result");
      expect(spy.mock.calls[0][0]).toBe("custom-span");

      const instrumented2 = instrumentation.instrument(() => "anonymous", {
        spanConfig: { spanName: "anonymous-span" },
      });

      expect(instrumented2()).toBe("anonymous");
      expect(spy.mock.calls[1][0]).toBe("anonymous-span");
    });

    test("should throw error when no span name available", () => {
      expect(() => instrumentation.instrument(() => "nothing")).toThrow();
    });

    test("should handle async functions correctly", async () => {
      const asyncFn = vi.fn().mockResolvedValue("async result");
      const instrumentedFn = instrumentation.instrument(asyncFn);

      const result = instrumentedFn();
      expect(result).toBeInstanceOf(Promise);

      const resolvedResult = await result;
      expect(resolvedResult).toBe("async result");
    });

    test("should handle sync errors and propagate them", () => {
      const error = new Error("sync error");
      const errorFn = vi.fn(() => {
        throw error;
      });

      const instrumentedFn = instrumentation.instrument(errorFn);
      expect(() => instrumentedFn()).toThrow(error);
    });

    test("should handle async errors and propagate them", async () => {
      const error = new Error("async error");
      const errorFn = vi.fn(async () => {
        await new Promise((resolve) => setTimeout(resolve, 10));
        throw error;
      });

      const errorFn2 = vi.fn(() => Promise.reject(error));

      // Non-catched error in an async function
      const instrumentedFn = instrumentation.instrument(errorFn);
      await expect(instrumentedFn()).rejects.toThrow(error);

      // Rejected promise in an async function
      const instrumentedFn2 = instrumentation.instrument(errorFn2);
      await expect(instrumentedFn2()).rejects.toThrow(error);
    });

    test("non-throwing functions should use isSuccessful predicate to determine success", () => {
      const fn = vi.fn(() => ({ data: "test", success: false }));
      const isSuccessful = vi.fn(
        (result: { success: boolean }) => result.success === true,
      );

      const instrumentedFn = instrumentation.instrument(fn, {
        isSuccessful,
        spanConfig: { spanName: "predicate-test" },
      });

      instrumentedFn();
      expect(isSuccessful).toHaveBeenCalledTimes(1);
    });

    test("should call onResult hook", () => {
      const onResult = vi.fn();
      const fn = vi.fn().mockReturnValue("result");

      const instrumentedFn = instrumentation.instrument(fn, {
        hooks: { onResult },
      });

      const nonSuccessfulInstrumentedFn = instrumentation.instrument(fn, {
        hooks: {
          onResult,
        },
        isSuccessful: () => false,
      });

      const result = instrumentedFn();
      const nonSuccessfulResult = nonSuccessfulInstrumentedFn();

      expect(result).toBe("result");
      expect(nonSuccessfulResult).toBe("result");

      expect(onResult).toHaveBeenCalledTimes(2);
    });

    test("should call onError hook", () => {
      const error = new Error("test error");
      const customError = new Error("custom error");
      const onError = vi.fn().mockReturnValue(customError);

      const fn = vi.fn(() => {
        throw error;
      });

      const instrumentedFn = instrumentation.instrument(fn, {
        hooks: { onError },
      });

      expect(() => instrumentedFn()).toThrow(error);
      expect(onError).toHaveBeenCalledOnce();
    });

    test("should handle non-Error exceptions", () => {
      const error = "string error";
      const fn = vi.fn(() => {
        throw error;
      });

      const instrumentedFn = instrumentation.instrument(fn, {
        spanConfig: { spanName: "non-error-test" },
      });

      expect(() => instrumentedFn()).toThrow(error);

      const namedInstrumentedFn = instrumentation.instrument(
        vi.fn(function named() {
          throw error;
        }),
        {},
      );

      expect(() => namedInstrumentedFn()).toThrow(error);
    });

    test("should pass through when telemetry is disabled", () => {
      vi.mocked(runtimeHelpers.isTelemetryEnabled).mockReturnValue(false);

      const fn = vi.fn((x: number) => {
        expect(() => instrumentation.getInstrumentationHelpers()).toThrow();
        return x * 2;
      });

      const instrumentedDuplicate = instrumentation.instrument(fn, {
        spanConfig: { spanName: "disabled-test" },
      });

      const result = instrumentedDuplicate(5);
      expect(result).toBe(10);
      expect(fn).toHaveBeenCalledWith(5);
    });

    test("should maintain separate contexts for concurrent operations", async () => {
      const contexts: InstrumentationContext[] = [];

      const fn1 = instrumentation.instrument(async function op1() {
        const ctx = instrumentation.getInstrumentationHelpers();
        await new Promise((resolve) => setTimeout(resolve, 15));
        contexts.push(ctx);

        return "op1";
      });

      const fn2 = instrumentation.instrument(async function op2() {
        const ctx = instrumentation.getInstrumentationHelpers();
        await new Promise((resolve) => setTimeout(resolve, 30));
        contexts.push(ctx);

        return "op2";
      });

      const [result1, result2] = await Promise.all([fn1(), fn2()]);
      expect(result1).toBe("op1");
      expect(result2).toBe("op2");

      const [ctx1, ctx2] = contexts;
      expect(ctx1).not.toBe(ctx2);
    });

    test("should propagate context through inner non-instrumented functions", () => {
      const contexts: InstrumentationContext[] = [];

      // Non-instrumented inner function
      const inner = vi.fn(function innerFn() {
        contexts.push(instrumentation.getInstrumentationHelpers());
        return "inner";
      });

      // Instrumented outer function
      const outer = instrumentation.instrument(function outerFn() {
        contexts.push(instrumentation.getInstrumentationHelpers());
        const result = inner();

        return `outer-${result}`;
      });

      const result = outer();
      expect(result).toBe("outer-inner");

      expect(contexts).toHaveLength(2);
      const [ctx1, ctx2] = contexts;

      expect(ctx1).toBe(ctx2);
      expect(ctx1.tracer).toBeDefined();
      expect(ctx2.tracer).toBeDefined();
      expect(ctx1.meter).toBeDefined();
      expect(ctx2.meter).toBeDefined();
    });

    test("inner instrumented functions should have their own context", () => {
      const contexts: InstrumentationContext[] = [];
      const inner = instrumentation.instrument(function innerFn() {
        contexts.push(instrumentation.getInstrumentationHelpers());
        return "inner";
      });

      const outer = instrumentation.instrument(function outerFn() {
        contexts.push(instrumentation.getInstrumentationHelpers());
        const result = inner();
        return `outer-${result}`;
      });

      expect(outer()).toBe("outer-inner");

      const [ctx1, ctx2] = contexts;
      expect(ctx1).not.toBe(ctx2);
    });
  });

  describe("instrumentEntrypoint", () => {
    let mockInitializeTelemetry: Mock<
      (
        params: Record<string, unknown>,
        isDevelopment: boolean,
      ) => TelemetryConfig
    >;
    let mockMain: Mock<
      (params: Record<string, unknown>) => { statusCode: number }
    >;

    beforeEach(() => {
      mockMain = vi.fn((_params: Record<string, unknown>) => ({
        statusCode: 200,
      }));

      mockInitializeTelemetry = vi.fn(
        (_params: Record<string, unknown>, _isDevelopment: boolean) => ({
          diagnostics: { logLevel: "error" },
          sdkConfig: {},
        }),
      );
    });

    test("should instrument entrypoint function and preserve behavior", async () => {
      const instrumentedMain = instrumentation.instrumentEntrypoint(mockMain, {
        initializeTelemetry: mockInitializeTelemetry,
      });

      const params = {};
      const result = await instrumentedMain(params);

      expect(result).toEqual({ statusCode: 200 });
      expect(mockMain).toHaveBeenCalledWith(params);
      expect(mockInitializeTelemetry).toHaveBeenCalled();
    });

    test.each([
      { ENABLE_TELEMETRY: "true", LOG_LEVEL: "debug" },
      { ENABLE_TELEMETRY: "false", LOG_LEVEL: "info" },
    ])("should always set telemetry environment variables", (params) => {
      const instrumentedMain = instrumentation.instrumentEntrypoint(mockMain, {
        initializeTelemetry: mockInitializeTelemetry,
      });

      instrumentedMain(params);
      expect(process.env.__AIO_LIB_TELEMETRY_LOG_LEVEL).toBe(params.LOG_LEVEL);
      expect(process.env.__AIO_LIB_TELEMETRY_ENABLE_TELEMETRY).toBe(
        params.ENABLE_TELEMETRY,
      );
    });

    test.each([
      { isDevelopment: true, logLevel: "debug" },
      { isDevelopment: false, logLevel: "info" },
    ])(
      "should default log level parameter when not provided depending on the environment",
      ({ isDevelopment, logLevel }) => {
        vi.mocked(runtimeHelpers.isDevelopment).mockReturnValue(isDevelopment);

        const instrumentedMain = instrumentation.instrumentEntrypoint(
          mockMain,
          {
            initializeTelemetry: mockInitializeTelemetry,
          },
        );

        instrumentedMain({});
        expect(process.env.__AIO_LIB_TELEMETRY_LOG_LEVEL).toBe(logLevel);
      },
    );

    test("should pass through when telemetry is disabled", async () => {
      // Temporarily override the telemetry enabled mock.
      vi.mocked(runtimeHelpers.isTelemetryEnabled).mockReturnValue(false);

      const instrumentedMain = instrumentation.instrumentEntrypoint(mockMain, {
        initializeTelemetry: mockInitializeTelemetry,
      });

      const params = {};
      const result = await instrumentedMain(params);

      expect(result).toEqual({ statusCode: 200 });
      expect(mockMain).toHaveBeenCalledWith(params);
      expect(mockInitializeTelemetry).not.toHaveBeenCalled();

      vi.mocked(runtimeHelpers.isTelemetryEnabled).mockReturnValue(true);
    });

    test.each([
      {
        name: "x-telemetry-context header",
        params: {
          __ow_headers: {
            "x-telemetry-context": JSON.stringify({
              traceparent:
                "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01",
            }),
          },
        },
      },
      {
        name: "__telemetryContext parameter",
        params: {
          __telemetryContext: {
            traceparent:
              "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01",
          },
        },
      },
      {
        name: "data.__telemetryContext parameter",
        params: {
          data: {
            __telemetryContext: {
              traceparent:
                "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01",
            },
          },
        },
      },
      {
        name: "W3C traceparent header",
        params: {
          __ow_headers: {
            traceparent:
              "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01",
          },
        },
      },
    ])("should handle context propagation from source: $name", ({ params }) => {
      let capturedSpanContext: SpanContext | null = null;
      const instrumentedMain = instrumentation.instrumentEntrypoint(
        function main() {
          capturedSpanContext = instrumentation
            .getInstrumentationHelpers()
            .currentSpan?.spanContext();

          return { statusCode: 200 };
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
        },
      );

      instrumentedMain({
        ENABLE_TELEMETRY: "true",
        ...params,
      });

      if (!capturedSpanContext) {
        expect.fail("capturedSpanContext is null");
      }

      const spanContext = capturedSpanContext as SpanContext;
      expect(spanContext.traceId).toBe("1234567890abcdef1234567890abcdef");
      expect(spanContext.spanId).not.toBe("1234567890abcdef"); // should create a new span
      expect(spanContext.traceFlags).toBe(1);
      expect(spanContext.traceState).toBe(undefined);
    });

    test("should skip propagation when configured", () => {
      let capturedSpanContext: SpanContext | null = null;
      const instrumentedMain = instrumentation.instrumentEntrypoint(
        function main() {
          capturedSpanContext = instrumentation
            .getInstrumentationHelpers()
            .currentSpan?.spanContext();

          return { statusCode: 200 };
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
          propagation: { skip: true },
        },
      );

      instrumentedMain({
        __telemetryContext: {
          traceparent:
            "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01",
        },
        ENABLE_TELEMETRY: "true",
      });

      if (!capturedSpanContext) {
        expect.fail("capturedContext is null");
      }

      const spanContext = capturedSpanContext as SpanContext;
      expect(spanContext).toBeDefined();
      expect(spanContext.traceId).not.toBe("1234567890abcdef1234567890abcdef");
    });

    test("should receive params in getContextCarrier", () => {
      let capturedParams: Record<string, unknown> | null = null;
      const instrumentedMain = instrumentation.instrumentEntrypoint(mockMain, {
        initializeTelemetry: mockInitializeTelemetry,
        propagation: {
          getContextCarrier: (entrypointParams) => {
            capturedParams = entrypointParams;
            return { baseCtx: undefined, carrier: {} };
          },
        },
      });

      const params = { ENABLE_TELEMETRY: "true" };
      instrumentedMain(params);

      expect(capturedParams).toBe(params);
    });

    test("should use custom base context if provided", () => {
      const baseContext = {
        deleteValue: vi.fn(),
        getValue: vi.fn(),
        setValue: vi.fn(),
      } satisfies Context;

      const spy = vi.spyOn(tracer, "startActiveSpan");
      const instrumentedMain = instrumentation.instrumentEntrypoint(
        function main() {
          return { statusCode: 200 };
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
          propagation: {
            getContextCarrier: () => ({ baseCtx: baseContext, carrier: {} }),
          },
        },
      );

      instrumentedMain({ ENABLE_TELEMETRY: "true" });
      expect(spy.mock.calls[0][2]).toEqual(baseContext);
    });

    test("should use default context if no base context is provided", () => {
      const spy = vi.spyOn(tracer, "startActiveSpan");
      const instrumentedMain = instrumentation.instrumentEntrypoint(
        function main() {
          return { statusCode: 200 };
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
          propagation: {
            getContextCarrier: () => ({ carrier: {} }),
          },
        },
      );

      instrumentedMain({ ENABLE_TELEMETRY: "true" });
      expect(spy.mock.calls[0][2]).toEqual(ROOT_CONTEXT);
    });

    test.each([new Error("Unexpected error"), "Unexpected string error"])(
      "should throw if there's some unexpected error during instrumentation",
      (error) => {
        vi.mocked(mockInitializeTelemetry).mockImplementation(() => {
          throw error;
        });

        const instrumentedMain = instrumentation.instrumentEntrypoint(
          mockMain,
          {
            initializeTelemetry: mockInitializeTelemetry,
          },
        );

        expect(() => instrumentedMain({})).toThrow();
      },
    );

    test("should bubble up sync runtime errors", () => {
      const error = new Error("Runtime error");
      mockMain.mockImplementation(() => {
        throw error;
      });

      const instrumentedMain = instrumentation.instrumentEntrypoint(mockMain, {
        initializeTelemetry: mockInitializeTelemetry,
      });

      expect(() => instrumentedMain({})).toThrow(error);
    });

    test("should bubble up async runtime errors", async () => {
      const error = new Error("Runtime error");
      mockMain.mockRejectedValueOnce(error);

      const instrumentedMain = instrumentation.instrumentEntrypoint(mockMain, {
        initializeTelemetry: mockInitializeTelemetry,
      });

      await expect(instrumentedMain({})).rejects.toThrow(error);
    });

    test("should provide instrumentation context to entrypoint", () => {
      let capturedContext: InstrumentationContext | null = null;

      const mainWithInstrumentation = vi.fn(() => {
        capturedContext = instrumentation.getInstrumentationHelpers();
        return { statusCode: 200 };
      });

      const instrumentedMain = instrumentation.instrumentEntrypoint(
        mainWithInstrumentation,
        {
          initializeTelemetry: mockInitializeTelemetry,
        },
      );

      instrumentedMain({ ENABLE_TELEMETRY: "true" });

      if (!capturedContext) {
        expect.fail("capturedContext is null");
      }

      const instrumentationContext = capturedContext as InstrumentationContext;
      expect(instrumentationContext).toBeDefined();
      expect(instrumentationContext.currentSpan).toBeDefined();
      expect(instrumentationContext.logger).toBeDefined();
      expect(instrumentationContext.tracer).toBeDefined();
      expect(instrumentationContext.meter).toBeDefined();
    });

    test("should handle async entrypoints", async () => {
      const asyncMain = vi.fn().mockResolvedValue({ statusCode: 200 });

      const instrumentedMain = instrumentation.instrumentEntrypoint(asyncMain, {
        initializeTelemetry: mockInitializeTelemetry,
      });

      const result = instrumentedMain({ ENABLE_TELEMETRY: "true" });
      expect(result).toBeInstanceOf(Promise);

      await expect(result).resolves.toEqual({ statusCode: 200 });
    });

    test("should set correct span name for entrypoint", () => {
      // For unnamed functions the span name is set to actionName/entrypoint
      const spy = vi.spyOn(tracer, "startActiveSpan");
      const unnamedEntrypoint = instrumentation.instrumentEntrypoint(
        () => "dummy",
        {
          initializeTelemetry: mockInitializeTelemetry,
        },
      );

      unnamedEntrypoint({});
      expect(spy.mock.calls[0][0]).toBe("test-action/entrypoint");

      const spanConfigNamedEntrypoint = instrumentation.instrumentEntrypoint(
        () => "dummy",
        {
          initializeTelemetry: mockInitializeTelemetry,
          spanConfig: { spanName: "custom-span" },
        },
      );

      spanConfigNamedEntrypoint({});
      expect(spy.mock.calls[1][0]).toBe("custom-span");

      const namedEntrypoint = instrumentation.instrumentEntrypoint(
        function main() {
          return "dummy";
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
        },
      );

      namedEntrypoint({});
      expect(spy.mock.calls[2][0]).toBe("test-action/main");
    });

    test("should mark root span as error if the runtime action fails", () => {
      const error = {
        message: "Runtime action failed",
        statusCode: 500,
      };

      let spy: MockInstance | null = null;
      const instrumentedMain = instrumentation.instrumentEntrypoint(
        function main() {
          const span = instrumentation.getInstrumentationHelpers().currentSpan;
          spy = vi.spyOn(span, "setStatus");

          return { error };
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
        },
      );

      const result = instrumentedMain({});
      expect(result).toEqual({ error });

      if (!spy) {
        expect.fail("failed to spy on span.setStatus");
      }

      const mockSpy = spy as MockInstance;
      expect(mockSpy.mock.calls[0][0]).toEqual({
        code: SpanStatusCode.ERROR,
      });
    });

    test("should not instrusively mark root span if the runtime action returns a non-object", () => {
      let spy: MockInstance | null = null;
      const instrumentedMain = instrumentation.instrumentEntrypoint(
        function main() {
          const span = instrumentation.getInstrumentationHelpers().currentSpan;
          spy = vi.spyOn(span, "setStatus");

          return "test";
        },
        {
          initializeTelemetry: mockInitializeTelemetry,
        },
      );

      const result = instrumentedMain({});
      expect(result).toEqual("test");

      if (!spy) {
        expect.fail("failed to spy on span.setStatus");
      }

      const mockSpy = spy as MockInstance;
      expect(mockSpy.mock.calls[0][0]).not.toEqual(SpanStatusCode.ERROR);
    });
  });
});
