def test_pex_builder():
  # test w/ and w/o zipfile dists
  with nested(temporary_dir(), make_bdist('p1', zipped=True)) as (td, p1):
    write_pex(td, exe_main, dists=[p1])

    success_txt = os.path.join(td, 'success.txt')
    PEX(td).run(args=[success_txt])
    assert os.path.exists(success_txt)
    with open(success_txt) as fp:
      assert fp.read() == 'success'

  # test w/ and w/o zipfile dists
  with nested(temporary_dir(), temporary_dir(), make_bdist('p1', zipped=True)) as (
      td1, td2, p1):
    target_egg_dir = os.path.join(td2, os.path.basename(p1.location))
    safe_mkdir(target_egg_dir)
    with closing(zipfile.ZipFile(p1.location, 'r')) as zf:
      zf.extractall(target_egg_dir)
    p1 = DistributionHelper.distribution_from_path(target_egg_dir)

    write_pex(td1, exe_main, dists=[p1])

    success_txt = os.path.join(td1, 'success.txt')
    PEX(td1).run(args=[success_txt])
    assert os.path.exists(success_txt)
    with open(success_txt) as fp:
      assert fp.read() == 'success'

def test_nested_requirements():
  with temporary_dir() as td1:
    with temporary_dir() as td2:
      with open(os.path.join(td1, 'requirements.txt'), 'w') as fp:
        fp.write(dedent('''
            requirement1
            requirement2
            -r %s
            -r %s
        ''' % (
            os.path.join(td2, 'requirements_nonrelative.txt'),
            os.path.join('relative', 'requirements_relative.txt'))
        ))

      with open(os.path.join(td2, 'requirements_nonrelative.txt'), 'w') as fp:
        fp.write(dedent('''
        requirement3
        requirement4
        '''))

      os.mkdir(os.path.join(td1, 'relative'))
      with open(os.path.join(td1, 'relative', 'requirements_relative.txt'), 'w') as fp:
        fp.write(dedent('''
        requirement5
        requirement6
        '''))

      def rr(req):
        return ResolvableRequirement.from_string(req, ResolverOptionsBuilder())

      reqs = requirements_from_file(os.path.join(td1, 'requirements.txt'))
      assert reqs == [rr('requirement%d' % k) for k in (1, 2, 3, 4, 5, 6)]

def test_source_packages():
  for ext in ('.tar.gz', '.tar', '.tgz', '.zip', '.tar.bz2'):
    sl = SourcePackage('a_p_r-3.1.3' + ext)
    assert sl._name == 'a_p_r'
    assert sl.name == 'a-p-r'
    assert sl.raw_version == '3.1.3'
    assert sl.version == parse_version(sl.raw_version)
    for req in ('a_p_r', 'a_p_r>2', 'a_p_r>3', 'a_p_r>=3.1.3', 'a_p_r==3.1.3', 'a_p_r>3,<3.5'):
      assert sl.satisfies(req)
      assert sl.satisfies(Requirement.parse(req))
    for req in ('foo', 'a_p_r==4.0.0', 'a_p_r>4.0.0', 'a_p_r>3.0.0,<3.0.3', 'a==3.1.3'):
      assert not sl.satisfies(req)
  sl = SourcePackage('python-dateutil-1.5.tar.gz')
  assert sl.name == 'python-dateutil'
  assert sl.raw_version == '1.5'

  with temporary_dir() as td:
    dateutil_base = 'python-dateutil-1.5'
    dateutil = '%s.zip' % dateutil_base
    with contextlib.closing(ZipFile(os.path.join(td, dateutil), 'w')) as zf:
      zf.writestr(os.path.join(dateutil_base, 'file1.txt'), 'junk1')
      zf.writestr(os.path.join(dateutil_base, 'file2.txt'), 'junk2')
    sl = SourcePackage('file://' + os.path.join(td, dateutil), opener=Web())
    with temporary_dir() as td2:
      sl.fetch(location=td2)
      print(os.listdir(td2))
      assert set(os.listdir(os.path.join(td2, dateutil_base))) == set(['file1.txt', 'file2.txt'])

def test_pex_builder_compilation():
  with nested(temporary_dir(), temporary_dir(), temporary_dir()) as (td1, td2, td3):
    src = os.path.join(td1, 'src.py')
    with open(src, 'w') as fp:
      fp.write(exe_main)

    exe = os.path.join(td1, 'exe.py')
    with open(exe, 'w') as fp:
      fp.write(exe_main)

    def build_and_check(path, precompile):
      pb = PEXBuilder(path)
      pb.add_source(src, 'lib/src.py')
      pb.set_executable(exe, 'exe.py')
      pb.freeze(bytecode_compile=precompile)
      for pyc_file in ('exe.pyc', 'lib/src.pyc', '__main__.pyc'):
        pyc_exists = os.path.exists(os.path.join(path, pyc_file))
        if precompile:
          assert pyc_exists
        else:
          assert not pyc_exists
      bootstrap_dir = os.path.join(path, PEXBuilder.BOOTSTRAP_DIR)
      bootstrap_pycs = []
      for _, _, files in os.walk(bootstrap_dir):
        bootstrap_pycs.extend(f for f in files if f.endswith('.pyc'))
      if precompile:
        assert len(bootstrap_pycs) > 0
      else:
        assert 0 == len(bootstrap_pycs)

    build_and_check(td2, False)
    build_and_check(td3, True)

def test_access_zipped_assets_integration():
  test_executable = dedent('''
      import os
      from _pex.util import DistributionHelper
      temp_dir = DistributionHelper.access_zipped_assets('my_package', 'submodule')
      with open(os.path.join(temp_dir, 'mod.py'), 'r') as fp:
        for line in fp:
          print(line)
  ''')
  with nested(temporary_dir(), temporary_dir()) as (td1, td2):
    pb = PEXBuilder(path=td1)
    with open(os.path.join(td1, 'exe.py'), 'w') as fp:
      fp.write(test_executable)
      pb.set_executable(fp.name)

    submodule = os.path.join(td1, 'my_package', 'submodule')
    safe_mkdir(submodule)
    mod_path = os.path.join(submodule, 'mod.py')
    with open(mod_path, 'w') as fp:
      fp.write('accessed')
      pb.add_source(fp.name, 'my_package/submodule/mod.py')

    pex = os.path.join(td2, 'app.pex')
    pb.build(pex)

    output, returncode = run_simple_pex(pex)
    try:
      output = output.decode('UTF-8')
    except ValueError:
      pass
    assert output == 'accessed\n'
    assert returncode == 0

def test_resolve_prereleases_cached():
  stable_dep = make_sdist(name='dep', version='2.0.0')
  prerelease_dep = make_sdist(name='dep', version='3.0.0rc3')

  with temporary_dir() as td:
    for sdist in (stable_dep, prerelease_dep):
      safe_copy(sdist, os.path.join(td, os.path.basename(sdist)))
    fetchers = [Fetcher([td])]

    with temporary_dir() as cd:
      def assert_resolve(dep, expected_version, **resolve_kwargs):
        dists = list(
          resolve_multi([dep], cache=cd, cache_ttl=1000, **resolve_kwargs)
        )
        assert 1 == len(dists)
        dist = dists[0]
        assert expected_version == dist.version

      Crawler.reset_cache()

      # First do a run to load it into the cache.
      assert_resolve('dep>=1,<4', '3.0.0rc3', allow_prereleases=True, fetchers=fetchers)

      # This simulates running from another pex command. The Crawler cache actually caches an empty
      # cache so this fails in the same "process".
      Crawler.reset_cache()

      # Now assert that we can get it from the cache by removing the source.
      assert_resolve('dep>=1,<4', '3.0.0rc3', allow_prereleases=True, fetchers=[])

      # It should also be able to resolve without allow_prereleases, if explicitly requested.
      Crawler.reset_cache()
      assert_resolve('dep>=1.rc1,<4', '3.0.0rc3', fetchers=[])

    def test_import_export(self):
        with temporary_dir() as state_dir:
            with temporary_dir() as server_dir1:
                with TestServer(config=dict(serverdir=server_dir1)) as devpi:
                    self.assertEqual(200, requests.get(devpi.url).status_code)
                export_state(server_dir1, state_dir)

            with temporary_dir() as server_dir2:
                import_state(server_dir2, state_dir)
                with TestServer(config=dict(serverdir=server_dir2)) as devpi:
                    self.assertEqual(200, requests.get(devpi.url).status_code)

def test_nested_pushd():
  pre_cwd = os.getcwd()
  with temporary_dir() as tempdir1:
    with pushd(tempdir1) as path1:
      assert os.getcwd() == os.path.realpath(tempdir1)
      with temporary_dir(root_dir=tempdir1) as tempdir2:
        with pushd(tempdir2) as path2:
          assert os.getcwd() == os.path.realpath(tempdir2)
        assert os.getcwd() == os.path.realpath(tempdir1)
      assert os.getcwd() == os.path.realpath(tempdir1)
    assert os.getcwd() == pre_cwd
  assert os.getcwd() == pre_cwd

def test_diamond_local_resolve_cached():
  # This exercises the issue described here: https://github.com/pantsbuild/pex/issues/120
  project1_sdist = make_sdist(name='project1', install_reqs=['project2<1.0.0'])
  project2_sdist = make_sdist(name='project2')

  with temporary_dir() as dd:
    for sdist in (project1_sdist, project2_sdist):
      safe_copy(sdist, os.path.join(dd, os.path.basename(sdist)))
    fetchers = [Fetcher([dd])]
    with temporary_dir() as cd:
      dists = resolve(['project1', 'project2'], fetchers=fetchers, cache=cd, cache_ttl=1000)
      assert len(dists) == 2

  def test_find_root_thrifts(self):
    with temporary_dir() as dir:
      root_1 = self.write(os.path.join(dir, 'root_1.thrift'), '# noop')
      root_2 = self.write(os.path.join(dir, 'root_2.thrift'), '# noop')
      self.assertEquals(set([root_1, root_2]),
                        find_root_thrifts(basedirs=[], sources=[root_1, root_2]))

    with temporary_dir() as dir:
      root_1 = self.write(os.path.join(dir, 'root_1.thrift'), 'include "mid_1.thrift"')
      self.write(os.path.join(dir, 'mid_1.thrift'), 'include "leaf_1.thrift"')
      self.write(os.path.join(dir, 'leaf_1.thrift'), '# noop')
      root_2 = self.write(os.path.join(dir, 'root_2.thrift'), 'include "root_1.thrift"')
      self.assertEquals(set([root_2]), find_root_thrifts(basedirs=[], sources=[root_1, root_2]))

 def test_round_trip(prefix=None):
   with temporary_dir() as fromdir:
     safe_mkdir(os.path.join(fromdir, 'a/b/c'))
     touch(os.path.join(fromdir, 'a/b/d/e.txt'))
     with temporary_dir() as archivedir:
       archive = archiver.create(fromdir, archivedir, 'archive', prefix=prefix)
       with temporary_dir() as todir:
         archiver.extract(archive, todir)
         fromlisting = listtree(fromdir)
         if prefix:
           fromlisting = set(os.path.join(prefix, x) for x in fromlisting)
           if empty_dirs:
             fromlisting.add(prefix)
         self.assertEqual(fromlisting, listtree(todir))

def test_pex_root():
  with temporary_dir() as tmp_home:
    with environment_as(HOME=tmp_home):
      with temporary_dir() as td:
        with temporary_dir() as output_dir:
          env = os.environ.copy()
          env['PEX_INTERPRETER'] = '1'

          output_path = os.path.join(output_dir, 'pex.pex')
          args = ['pex', '-o', output_path, '--not-zip-safe', '--pex-root={0}'.format(td)]
          results = run_pex_command(args=args, env=env)
          results.assert_success()
          assert ['pex.pex'] == os.listdir(output_dir), 'Expected built pex file.'
          assert [] == os.listdir(tmp_home), 'Expected empty temp home dir.'
          assert 'build' in os.listdir(td), 'Expected build directory in tmp pex root.'

  def test_sibling_references(self):
    with temporary_dir() as root_dir:
      buildfile = create_buildfile(root_dir, 'a', name='BUILD',
        content=dedent("""
          dependencies(name='util',
            dependencies=[
              jar(org='com.twitter', name='util', rev='0.0.1')
            ]
          )
        """).strip()
      )
      sibling = create_buildfile(root_dir, 'a', name='BUILD.sibling',
        content=dedent("""
          dependencies(name='util-ex',
            dependencies=[
              pants(':util'),
              jar(org='com.twitter', name='util-ex', rev='0.0.1')
            ]
          )
        """).strip()
      )
      ParseContext(buildfile).parse()

      utilex = Target.get(Address.parse(root_dir, 'a:util-ex', is_relative=False))
      utilex_deps = set(utilex.resolve())

      util = Target.get(Address.parse(root_dir, 'a:util', is_relative=False))
      util_deps = set(util.resolve())

      self.assertEquals(util_deps, util_deps.intersection(utilex_deps))

def test_empty_resolve():
  empty_resolve = resolve([])
  assert empty_resolve == []

  with temporary_dir() as td:
    empty_resolve = resolve([], cache=td)
    assert empty_resolve == []

  def test_filesystem_image_containerizer_not_executable(self):
    proxy_driver = ProxyDriver()

    with temporary_dir() as tempdir:

      tempfile = os.path.join(tempdir, 'fake-containierizer')
      with open(tempfile, 'a'):
        os.utime(tempfile, None)

      te = FastThermosExecutor(
        runner_provider=make_provider(tempdir, mesos_containerizer_path=tempfile),
        sandbox_provider=FileSystemImageTestSandboxProvider())

      te.SANDBOX_INITIALIZATION_TIMEOUT = Amount(1, Time.MILLISECONDS)
      te.START_TIMEOUT = Amount(10, Time.MILLISECONDS)
      te.STOP_TIMEOUT = Amount(10, Time.MILLISECONDS)

      te.launchTask(proxy_driver, make_task(HELLO_WORLD_MTI))

      proxy_driver.wait_stopped()

      updates = proxy_driver.method_calls['sendStatusUpdate']
      assert len(updates) == 2
      assert updates[0][0][0].state == mesos_pb2.TASK_STARTING
      assert updates[1][0][0].state == mesos_pb2.TASK_FAILED

  def test_killTask_during_runner_initialize(self):  # noqa
    proxy_driver = ProxyDriver()

    task = make_task(HELLO_WORLD_MTI)

    with temporary_dir() as td:
      te = FastThermosExecutor(
          runner_provider=make_provider(td),
          sandbox_provider=SlowSandboxProvider())
      te.launchTask(proxy_driver, task)
      te.sandbox_initialized.wait()
      te.killTask(proxy_driver, mesos_pb2.TaskID(value=task.task_id.value))
      assert te.runner_aborted.is_set()
      assert not te.sandbox_created.is_set()

      # we've simulated a "slow" initialization by blocking it until the killTask was sent - so now,
      # trigger the initialization to complete
      te._sandbox._init_start.set()

      # however, wait on the runner to definitely finish its initialization before continuing
      # (otherwise, this function races ahead too fast)
      te._sandbox._init_done.wait()
      te.sandbox_created.wait()
      assert te.sandbox_initialized.is_set()
      assert te.sandbox_created.is_set()

      proxy_driver.wait_stopped()

      updates = proxy_driver.method_calls['sendStatusUpdate']
      assert len(updates) == 2
      assert updates[-1][0][0].state == mesos_pb2.TASK_KILLED

  def relativize_analysis_file(self, src, dst):
    # Make an analysis cache portable. Work on a tmpfile, for safety.
    #
    # NOTE: We can't port references to deps on the Java home. This is because different JVM
    # implementations on different systems have different structures, and there's not
    # necessarily a 1-1 mapping between Java jars on different systems. Instead we simply
    # drop those references from the analysis file.
    #
    # In practice the JVM changes rarely, and it should be fine to require a full rebuild
    # in those rare cases.
    with temporary_dir() as tmp_analysis_dir:
      stripped_src = os.path.join(tmp_analysis_dir, 'analysis.nojava')
      tmp_analysis_file = os.path.join(tmp_analysis_dir, 'analysis.relativized')

      # Strip all lines containing self._java_home.
      with open(src, 'r') as infile:
        with open (stripped_src, 'w') as outfile:
          for line in infile:
            if not self._java_home in line:
              outfile.write(line)

      rebasings = [
        (self._ivy_home, ZincUtils.IVY_HOME_PLACEHOLDER),
        (self._pants_home, ZincUtils.PANTS_HOME_PLACEHOLDER),
      ]
      Analysis.rebase(stripped_src, tmp_analysis_file, rebasings)
      shutil.move(tmp_analysis_file, dst)

def test_verify_user_match(mock_check_output):
    with temporary_dir() as d:
        sandbox = FileSystemImageSandbox(os.path.join(d, "sandbox"), user="someuser", sandbox_mount_point="/some/path")

        mock_check_output.return_value = "uid=1(test-user) gid=2(test-group) groups=2(test-group)"
        # valid case
        sandbox._verify_user_match_in_taskfs(1, "test-user", 2, "test-group")
        mock_check_output.assert_called_with(["chroot", sandbox._task_fs_root, "id", "test-user"])

        # invalid user id
        with pytest.raises(FileSystemImageSandbox.CreationError):
            sandbox._verify_user_match_in_taskfs(0, "test-user", 2, "test-group")

        # invalid user name
        with pytest.raises(FileSystemImageSandbox.CreationError):
            sandbox._verify_user_match_in_taskfs(1, "invalid-user", 2, "test-group")

        # invalid group id
        with pytest.raises(FileSystemImageSandbox.CreationError):
            sandbox._verify_user_match_in_taskfs(1, "test-user", 0, "test-group")

        # invalid group name
        with pytest.raises(FileSystemImageSandbox.CreationError):
            sandbox._verify_user_match_in_taskfs(1, "test-user", 2, "invalid-group")

        # exception case
        exception = subprocess.CalledProcessError(returncode=1, cmd="some command", output=None)
        mock_check_output.side_effect = exception
        with pytest.raises(FileSystemImageSandbox.CreationError):
            sandbox._verify_user_match_in_taskfs(1, "test-user", 2, "test-group")

def test_verify_group_match(mock_check_output):
  with temporary_dir() as d:
    sandbox = FileSystemImageSandbox(d, user='someuser', sandbox_mount_point='/some/path')

    mock_check_output.return_value = 'test-group:x:2:'

    # valid case
    sandbox._verify_group_match_in_taskfs(2, 'test-group')
    mock_check_output.assert_called_with(
        ['chroot', sandbox._task_fs_root, 'getent', 'group', 'test-group'])

    # invalid group id
    with pytest.raises(FileSystemImageSandbox.CreationError):
      sandbox._verify_group_match_in_taskfs(3, 'test-group')

    # invalid group name
    with pytest.raises(FileSystemImageSandbox.CreationError):
      sandbox._verify_group_match_in_taskfs(2, 'invalid-group')

    # exception case
    exception = subprocess.CalledProcessError(
        returncode=1,
        cmd='some command',
        output=None)
    mock_check_output.side_effect = exception
    with pytest.raises(FileSystemImageSandbox.CreationError):
      sandbox._verify_group_match_in_taskfs(2, 'test-group')

def test_simple_process_filesystem_isolator():
    with temporary_dir() as td:
        taskpath = make_taskpath(td)
        sandbox = setup_sandbox(td, taskpath)

        test_isolator_path = os.path.join(td, "fake-mesos-containerier")
        with open(test_isolator_path, "w") as fd:
            # We use a fake version of the mesos-containerizer binary that just echoes out its args so
            # we can assert on them in the process's output.
            fd.write("\n".join(["#!/bin/sh", 'echo "[email protected]"']))

            fd.close()

            chmod_plus_x(test_isolator_path)

            p = TestProcess(
                "process", "echo hello world", 0, taskpath, sandbox, mesos_containerizer_path=test_isolator_path
            )
            p.start()

        rc = wait_for_rc(taskpath.getpath("process_checkpoint"))
        assert rc == 0
        assert_log_content(
            taskpath,
            "stdout",
            "launch --unshare_namespace_mnt --rootfs=/some/path/taskfs --user=None "
            '--command={"shell":true,"value":"echo hello world"}\n',
        )