Compile different function on Mix env and test

my_module.ex

defmodule MyModule do
  @moduledoc false

  if Mix.env() == :prod do
    @foo Application.fetch_env!(:myapp, :foo)
    def foo, do: @foo
  else
    def foo, do: Application.fetch_env!(:myapp, :foo)
  end
end

my_module_test.exs

defmodule MyModuleTest do
  use ExUnit.Case

  import ExUnit.CaptureIO

  @source_filepath __DIR__
                   |> Path.join("../../lib/myapp/my_module.ex")
                   |> Path.expand()
  test "foo" do
    on_exit(fn ->
      Mix.env(:test)
    end)

    # suppress warnings
    capture_io(:stderr, fn ->
      Mix.env(:prod)
      Code.compile_file(@source_filepath)
      assert MyModule.foo(), "compile on Mix.env(:prod)"

      Mix.env(:test)
      Code.compile_file(@source_filepath)
      assert MyModule.foo(), "compile on Mix.env(:test)"
    end)
  end
end