008 7of10

defmodule Foundation.TestSupport do
  @moduledoc "Utilities for testing Foundation-based applications."

  @doc """
  Executes a block of code with a service temporarily replaced by a mock.

  This macro ensures that any lookup for the specified `service_name` within
  the `do` block will resolve to the provided `mock_pid`. The mocking is
  isolated to the current test process and is safe for concurrent execution.

  ## Parameters
    - `service_tuple`: A `{namespace, service_name}` tuple identifying the service to mock.
    - `mock_pid`: The PID of the GenServer or process to use as the mock.
    - `block`: The block of code to execute within the mocked context.
  """
  defmacro with_mocked_service({namespace, service_name}, mock_pid, do: block) do
    quote do
      # 1. Get the unique test reference from the ExUnit context.
      test_ref = Foundation.TestSupport.get_test_ref()
      test_namespace = {:test, test_ref}
      
      # 2. Register the mock PID under the test's unique namespace.
      #    We register it with the *production* service name.
      Foundation.ServiceRegistry.register(test_namespace, unquote(service_name), unquote(mock_pid))
      
      # 3. Temporarily override the process dictionary to point to our test namespace.
      original_namespace = Foundation.ServiceRegistry.get_current_namespace()
      Foundation.ServiceRegistry.set_current_namespace(test_namespace)
      
      try
        # 4. Execute the user's test code.
        unquote(block)
      after
        # 5. Restore the original namespace and clean up.
        Foundation.ServiceRegistry.set_current_namespace(original_namespace)
        # The on_exit hook from setup_isolated_dspy_environment will handle the unregistering.
      end
    end
  end

  # Helper to get the test ref, assumes it's set in a setup block.
  def get_test_ref do
    Process.get(:current_test_ref) || raise "No test reference found. Did you use setup_isolated_dspy_environment()?"
  end
end

defmodule Foundation.ServiceRegistry do
  # ... existing functions ...
  
  @doc "Sets the active namespace for the calling process."
  def set_current_namespace(namespace) do
    Process.put(:foundation_current_namespace, namespace)
  end

  @doc "Gets the active namespace for the calling process, defaulting to :production."
  def get_current_namespace do
    Process.get(:foundation_current_namespace, :production)
  end
  
  # The `lookup` function must be modified to use this new process-local value.
  @spec lookup(service_name :: atom()) :: {:ok, pid()} | {:error, Error.t()}
  def lookup(service_name) do
    # This function no longer takes a namespace; it discovers it.
    namespace = get_current_namespace()
    
    # First, try to look up in the current (potentially test) namespace.
    case ProcessRegistry.lookup(namespace, service_name) do
      {:ok, pid} ->
        {:ok, pid}
      
      :error when namespace != :production ->
        # If in a test namespace and not found, fall back to production namespace.
        # This allows mocking only a subset of services while others remain real.
        ProcessRegistry.lookup(:production, service_name)

      :error ->
        # Not found in production namespace either.
        {:error, create_service_not_found_error(namespace, service_name)}
    end
  end
end

# in foundation/test_support.ex (or dspex/test/support.ex)

defmacro __using__(_opts) do
  quote do
    use ExUnit.Case, async: true
    import Foundation.TestSupport

    setup do
      test_ref = make_ref()
      Process.put(:current_test_ref, test_ref)
      
      # Optional: Start common services in the test namespace if needed.
      # start_supervised!({MyService, namespace: {:test, test_ref}})
      
      on_exit(fn ->
        # This is crucial for cleanup after each test.
        Foundation.ServiceRegistry.cleanup_test_namespace(test_ref)
      end)
      
      :ok
    end
  end
end

# Example usage in a test file
defmodule MyApp.MyTest do
  use MyApp.DataCase # This would 'use Foundation.TestSupport'
  
  test "..." do
    # ...
  end
end

# in test/dspex/predict_test.exs

defmodule DSPEx.PredictTest do
  use DSPEx.TestCase # Assumes this uses the Foundation.TestSupport setup

  # Define a simple mock LM client
  defmodule MockLM do
    use GenServer
    def start_link(opts), do: GenServer.start_link(__MODULE__, opts, [])
    def init(opts), do: {:ok, opts}
    
    def handle_call({:request, _messages, _opts}, _from, state) do
      # Return a canned response based on the test's needs
      response = Keyword.get(state, :response, {:ok, %{"choices" => [%{"message" => %{"content" => "mocked answer"}}]}})
      {:reply, response, state}
    end
  end

  test "forward returns ok tuple on successful LM call" do
    # 1. Start a mock LM GenServer configured for this specific test case.
    {:ok, mock_lm_pid} = MockLM.start_link(response: {:ok, %{"content" => "Test Passed"}})
    
    # 2. Use the macro to sandbox the execution.
    #    We mock the :openai_client service which is registered in the :production namespace.
    with_mocked_service({:production, :openai_client}, mock_lm_pid) do
      # 3. Inside this block, any code that calls `ServiceRegistry.lookup(:openai_client)`
      #    will receive `mock_lm_pid`.
      
      # The Predict module is configured to look up the :openai_client by name.
      predict_program = %DSPEx.Predict{signature: MyTestSignature, client: :openai_client}
      
      # Run the program.
      result = DSPEx.Program.forward(predict_program, %{question: "test"})
      
      assert {:ok, %DSPEx.Prediction{answer: "Test Passed"}} = result
    end
  end
end