#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file 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 CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Unit tests for the stager module."""

# pytype: skip-file

import io
import logging
import os
import shutil
import sys
import tempfile
import unittest

import mock
import pytest

from apache_beam.internal import pickler
from apache_beam.io.filesystem import CompressionTypes
from apache_beam.io.filesystems import FileSystems
from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.runners.internal import names
from apache_beam.runners.portability import stager

_LOGGER = logging.getLogger(__name__)

try:
  import setuptools
except ImportError:
  setuptools = None


class StagerTest(unittest.TestCase):
  def setUp(self):
    self._temp_dir = None
    self.stager = TestStager()

  def tearDown(self):
    if self._temp_dir:
      shutil.rmtree(self._temp_dir)
    self.stager = None
    # [https://github.com/apache/beam/issues/21457] set pickler to dill by
    # default.
    pickler.set_library(pickler.DEFAULT_PICKLE_LIB)

  def make_temp_dir(self):
    if self._temp_dir is None:
      self._temp_dir = tempfile.mkdtemp()
    return tempfile.mkdtemp(dir=self._temp_dir)

  def update_options(self, options):
    setup_options = options.view_as(SetupOptions)
    setup_options.sdk_location = ''

  def create_temp_file(self, path, contents):
    with open(path, 'w') as f:
      f.write(contents)
      return f.name

  # We can not rely on actual remote file systems paths hence making
  # '/tmp/remote/' a new remote path.
  def is_remote_path(self, path):
    return path.startswith('/tmp/remote/')

  remote_copied_files: list[str] = []

  def file_copy(self, from_path, to_path):
    if self.is_remote_path(from_path):
      self.remote_copied_files.append(from_path)
      _, from_name = os.path.split(from_path)
      if os.path.isdir(to_path):
        to_path = os.path.join(to_path, from_name)
      self.create_temp_file(to_path, 'nothing')
      _LOGGER.info('Fake copied remote file: %s to %s', from_path, to_path)
    elif self.is_remote_path(to_path):
      _LOGGER.info('Faking upload_file(%s, %s)', from_path, to_path)
    else:
      shutil.copyfile(from_path, to_path)

  def populate_requirements_cache(
      self, requirements_file, cache_dir, populate_cache_with_sdists=False):
    _ = requirements_file
    self.create_temp_file(os.path.join(cache_dir, 'abc.txt'), 'nothing')
    self.create_temp_file(os.path.join(cache_dir, 'def.txt'), 'nothing')
    # Return the list of packages that were "downloaded" for this requirements
    return {'abc.txt', 'def.txt'}

  @mock.patch('apache_beam.runners.portability.stager.open')
  @mock.patch('apache_beam.runners.portability.stager.get_new_http')
  def test_download_file_https(self, mock_new_http, mock_open):
    from_url = 'https://from_url'
    to_path = '/tmp/http_file/'
    mock_new_http.return_value.request.return_value = ({
        'status': 200
    },
                                                       'file_content')
    self.stager._download_file(from_url, to_path)
    assert mock_open.mock_calls == [
        mock.call('/tmp/http_file/', 'wb'),
        mock.call().__enter__(),
        mock.call().__enter__().write('file_content'),
        mock.call().__exit__(None, None, None)
    ]

  @mock.patch('apache_beam.runners.portability.stager.open')
  @mock.patch('apache_beam.runners.portability.stager.get_new_http')
  @mock.patch.object(FileSystems, 'open')
  def test_download_file_non_http(self, mock_fs_open, mock_new_http, mock_open):
    from_url = 'gs://bucket/from_url'
    to_path = '/tmp/file/'
    mock_fs_open.return_value = io.BytesIO(b"file_content")
    self.stager._download_file(from_url, to_path)
    assert not mock_new_http.called
    mock_fs_open.assert_called_with(
        from_url, compression_type=CompressionTypes.UNCOMPRESSED)
    assert mock_open.mock_calls == [
        mock.call('/tmp/file/', 'wb'),
        mock.call().__enter__(),
        mock.call().__enter__().write(b'file_content'),
        mock.call().__exit__(None, None, None)
    ]

  @mock.patch('apache_beam.runners.portability.stager.os.mkdir')
  @mock.patch('apache_beam.runners.portability.stager.shutil.copyfile')
  @mock.patch('apache_beam.runners.portability.stager.get_new_http')
  def test_download_file_unrecognized(
      self, mock_new_http, mock_copyfile, mock_mkdir):
    from_url = '/tmp/from_file'
    to_path = '/tmp/to_file/'
    with mock.patch('apache_beam.runners.portability.stager.os.path.isdir',
                    return_value=True):
      self.stager._download_file(from_url, to_path)
      assert not mock_new_http.called
      mock_copyfile.assert_called_with(from_url, to_path)

    with mock.patch('apache_beam.runners.portability.stager.os.path.isdir',
                    return_value=False):
      self.stager._download_file(from_url, to_path)
      assert mock_mkdir.called

  def test_no_staging_location(self):
    with self.assertRaises(RuntimeError) as cm:
      self.stager.stage_job_resources([], staging_location=None)
    self.assertEqual(
        'The staging_location must be specified.', cm.exception.args[0])

  def test_no_main_session(self):
    staging_dir = self.make_temp_dir()
    options = PipelineOptions()

    options.view_as(SetupOptions).save_main_session = False
    self.update_options(options)

    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])

  # xdist adds unpicklable modules to the main session.
  @pytest.mark.no_xdist
  @pytest.mark.uses_dill
  @unittest.skipIf(
      sys.platform == "win32" and sys.version_info < (3, 8),
      'https://github.com/apache/beam/issues/20659: pytest on Windows pulls '
      'in a zipimporter, unpicklable before py3.8')
  def test_with_main_session(self):
    pytest.importorskip("dill")
    staging_dir = self.make_temp_dir()
    options = PipelineOptions()

    options.view_as(SetupOptions).save_main_session = True
    options.view_as(SetupOptions).pickle_library = pickler.USE_DILL
    self.update_options(options)

    self.assertEqual([
        names.PICKLED_MAIN_SESSION_FILE,
        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
    ],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])
    self.assertTrue(
        os.path.isfile(
            os.path.join(staging_dir, names.PICKLED_MAIN_SESSION_FILE)))

  # (https://github.com/apache/beam/issues/21457): Remove the decorator once
  # cloudpickle is default pickle library
  @pytest.mark.no_xdist
  def test_main_session_staged_when_using_cloudpickle(self):
    staging_dir = self.make_temp_dir()
    options = PipelineOptions()

    options.view_as(SetupOptions).save_main_session = True
    # even if the save main session is on, no pickle file for main
    # session is saved when pickle_library==cloudpickle.
    options.view_as(SetupOptions).pickle_library = pickler.USE_CLOUDPICKLE
    self.update_options(options)
    self.assertEqual([
        names.PICKLED_MAIN_SESSION_FILE,
        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
    ],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])

  def test_default_resources(self):
    staging_dir = self.make_temp_dir()
    options = PipelineOptions()
    self.update_options(options)

    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])

  def test_no_submission_env_staging(self):
    staging_dir = self.make_temp_dir()
    options = PipelineOptions()
    self.update_options(options)

    resources = self.stager.create_job_resources(
        options, staging_dir, log_submission_env_dependencies=False)
    self.assertEqual([], resources)

  def test_with_requirements_file(self):
    staging_dir = self.make_temp_dir()
    requirements_cache_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).requirements_cache = requirements_cache_dir
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
    self.assertEqual(
        sorted([
            stager.REQUIREMENTS_FILE,
            'abc.txt',
            'def.txt',
            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
        ]),
        sorted(
            self.stager.create_and_stage_job_resources(
                options,
                populate_requirements_cache=self.populate_requirements_cache,
                staging_location=staging_dir)[1]))
    self.assertTrue(
        os.path.isfile(os.path.join(staging_dir, stager.REQUIREMENTS_FILE)))

  def test_with_pypi_requirements(self):
    staging_dir = self.make_temp_dir()
    requirements_cache_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).requirements_cache = requirements_cache_dir
    resources = self.stager.create_and_stage_job_resources(
        options,
        pypi_requirements=['nothing>=1.0,<2.0'],
        populate_requirements_cache=self.populate_requirements_cache,
        staging_location=staging_dir)[1]
    self.assertEqual(4, len(resources))
    self.assertTrue({'abc.txt', 'def.txt'} <= set(resources))
    generated_requirements = (
        set(resources) -
        {'abc.txt', 'def.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE}).pop()
    with open(os.path.join(staging_dir, generated_requirements)) as f:
      data = f.read()
    self.assertEqual('nothing>=1.0,<2.0', data)
    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt')))

  def test_requirements_file_not_present(self):
    staging_dir = self.make_temp_dir()
    with self.assertRaises(RuntimeError) as cm:
      options = PipelineOptions()
      self.update_options(options)
      options.view_as(SetupOptions).requirements_file = 'nosuchfile'
      self.stager.create_and_stage_job_resources(
          options,
          populate_requirements_cache=self.populate_requirements_cache,
          staging_location=staging_dir)
    self.assertEqual(
        cm.exception.args[0],
        'The file %s cannot be found. It was specified in the '
        '--requirements_file command line option.' % 'nosuchfile')

  def test_with_requirements_file_and_cache(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    options.view_as(SetupOptions).requirements_cache = self.make_temp_dir()
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
    self.assertEqual(
        sorted([
            stager.REQUIREMENTS_FILE,
            'abc.txt',
            'def.txt',
            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
        ]),
        sorted(
            self.stager.create_and_stage_job_resources(
                options,
                populate_requirements_cache=self.populate_requirements_cache,
                staging_location=staging_dir)[1]))
    self.assertTrue(
        os.path.isfile(os.path.join(staging_dir, stager.REQUIREMENTS_FILE)))
    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt')))

  def test_requirements_cache_not_populated_when_cache_disabled(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    options.view_as(
        SetupOptions).requirements_cache = stager.SKIP_REQUIREMENTS_CACHE
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
    with mock.patch(
        'apache_beam.runners.portability.stager_test.StagerTest.'
        'populate_requirements_cache') as (populate_requirements_cache):
      resources = self.stager.create_and_stage_job_resources(
          options,
          populate_requirements_cache=self.populate_requirements_cache,
          staging_location=staging_dir)[1]
      assert not populate_requirements_cache.called
      self.assertEqual(
          [stager.REQUIREMENTS_FILE, stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
          resources)
      self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
      self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'def.txt')))

  def test_with_pypi_requirements_skipping_cache(self):
    staging_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(
        SetupOptions).requirements_cache = stager.SKIP_REQUIREMENTS_CACHE

    resources = self.stager.create_and_stage_job_resources(
        options,
        pypi_requirements=['nothing>=1.0,<2.0'],
        populate_requirements_cache=self.populate_requirements_cache,
        staging_location=staging_dir)[1]
    with open(os.path.join(staging_dir, resources[0])) as f:
      data = f.read()
    self.assertEqual('nothing>=1.0,<2.0', data)
    self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
    self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'def.txt')))

  def test_setup_file_not_present(self):
    staging_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).setup_file = 'nosuchfile'

    with self.assertRaises(RuntimeError) as cm:
      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertEqual(
        cm.exception.args[0],
        'The file %s cannot be found. It was specified in the '
        '--setup_file command line option.' % 'nosuchfile')

  def test_setup_file_not_named_setup_dot_py(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).setup_file = (
        os.path.join(source_dir, 'xyz-setup.py'))

    self.create_temp_file(os.path.join(source_dir, 'xyz-setup.py'), 'notused')
    with self.assertRaises(RuntimeError) as cm:
      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertTrue(
        cm.exception.args[0].startswith(
            'The --setup_file option expects the full path to a file named '
            'setup.py or pyproject.toml instead of '))

  def test_setup_file_supplies_unexpected_filename(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).setup_file = (
        os.path.join(source_dir, 'xyz-pyproject.toml'))

    self.create_temp_file(
        os.path.join(source_dir, 'xyz-pyproject.toml'), 'notused')
    with self.assertRaises(RuntimeError) as cm:
      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertTrue(
        cm.exception.args[0].startswith(
            'The --setup_file option expects the full path to a file named '
            'setup.py or pyproject.toml instead of '))

  @unittest.skipIf(setuptools is None, "setuptools not available in this env")
  def test_setup_file_as_setup_dot_py_file(self):
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).setup_file = (
        os.path.join(source_dir, 'setup.py'))

    self.create_temp_file(
        os.path.join(source_dir, 'setup.py'),
        "from setuptools import setup; setup(name='my_package')")
    temp_dir = tempfile.mkdtemp()
    resources = self.stager.create_job_resources(
        options=options, temp_dir=temp_dir)
    file_path = list(self.stager.extract_staging_tuple_iter(resources))[0][0]

    self.assertEqual(
        os.path.join(temp_dir, 'my_package-0.0.0.tar.gz'), file_path)

  def test_sdk_location_default(self):
    staging_dir = self.make_temp_dir()
    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = 'default'

    _, staged_resources = self.stager.create_and_stage_job_resources(
        options, temp_dir=self.make_temp_dir(), staging_location=staging_dir)

    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                     staged_resources)

  def test_sdk_location_local_directory(self):
    staging_dir = self.make_temp_dir()
    sdk_location = self.make_temp_dir()
    self.create_temp_file(
        os.path.join(sdk_location, names.STAGED_SDK_SOURCES_FILENAME),
        'Package content.')

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = sdk_location

    self.assertEqual([
        names.STAGED_SDK_SOURCES_FILENAME,
        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
    ],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])
    tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME)
    with open(tarball_path) as f:
      self.assertEqual(f.read(), 'Package content.')

  def test_files_to_stage(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    foo_ca = os.path.join(source_dir, 'foo.ca')
    self.create_temp_file(foo_ca, 'ca content')
    test_txt = os.path.join(source_dir, 'test.txt')
    self.create_temp_file(test_txt, 'test content')
    files_to_stage = [foo_ca, test_txt]
    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).files_to_stage = files_to_stage

    self.assertEqual(
        ['foo.ca', 'test.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
        self.stager.create_and_stage_job_resources(
            options, staging_location=staging_dir)[1])
    foo = os.path.join(staging_dir, 'foo.ca')
    with open(foo) as f:
      self.assertEqual(f.read(), 'ca content')
    txt = os.path.join(staging_dir, 'test.txt')
    with open(txt) as f:
      self.assertEqual(f.read(), 'test content')

  def test_sdk_location_local_source_file(self):
    staging_dir = self.make_temp_dir()
    sdk_directory = self.make_temp_dir()
    sdk_filename = 'apache-beam-3.0.0.tar.gz'
    sdk_location = os.path.join(sdk_directory, sdk_filename)
    self.create_temp_file(sdk_location, 'Package content.')

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = sdk_location

    self.assertEqual([
        names.STAGED_SDK_SOURCES_FILENAME,
        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
    ],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])
    tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME)
    with open(tarball_path) as f:
      self.assertEqual(f.read(), 'Package content.')

  def test_sdk_location_local_wheel_file(self):
    staging_dir = self.make_temp_dir()
    sdk_directory = self.make_temp_dir()
    sdk_filename = 'apache_beam-1.0.0-cp27-cp27mu-manylinux1_x86_64.whl'
    sdk_location = os.path.join(sdk_directory, sdk_filename)
    self.create_temp_file(sdk_location, 'Package content.')

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = sdk_location

    self.assertEqual([sdk_filename, stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])
    tarball_path = os.path.join(staging_dir, sdk_filename)
    with open(tarball_path) as f:
      self.assertEqual(f.read(), 'Package content.')

  def test_sdk_location_local_directory_not_present(self):
    staging_dir = self.make_temp_dir()
    sdk_location = 'nosuchdir'
    with self.assertRaises(RuntimeError) as cm:
      options = PipelineOptions()
      self.update_options(options)
      options.view_as(SetupOptions).sdk_location = sdk_location

      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertEqual(
        'The file "%s" cannot be found. Its '
        'location was specified by the --sdk_location command-line option.' %
        sdk_location,
        cm.exception.args[0])

  @mock.patch(
      'apache_beam.runners.portability.stager_test.TestStager.stage_artifact')
  @mock.patch(
      'apache_beam.runners.portability.stager_test.stager.Stager._download_file'
  )
  def test_sdk_location_remote_source_file(self, *unused_mocks):
    staging_dir = self.make_temp_dir()
    sdk_location = 'gs://my-gcs-bucket/tarball.tar.gz'

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = sdk_location

    self.assertEqual([
        names.STAGED_SDK_SOURCES_FILENAME,
        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
    ],
                     self.stager.create_and_stage_job_resources(
                         options, staging_location=staging_dir)[1])

  def test_sdk_location_remote_wheel_file(self, *unused_mocks):
    staging_dir = self.make_temp_dir()
    sdk_filename = 'apache_beam-1.0.0-cp27-cp27mu-manylinux1_x86_64.whl'
    sdk_location = 'https://storage.googleapis.com/my-gcs-bucket/' + \
                   sdk_filename

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = sdk_location

    def file_download(_, to_path):
      with open(to_path, 'w') as f:
        f.write('Package content.')
      return to_path

    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._download_file',
                    staticmethod(file_download)):
      self.assertEqual([sdk_filename, stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                       self.stager.create_and_stage_job_resources(
                           options, staging_location=staging_dir)[1])

    wheel_file_path = os.path.join(staging_dir, sdk_filename)
    with open(wheel_file_path) as f:
      self.assertEqual(f.read(), 'Package content.')

  def test_sdk_location_http(self):
    staging_dir = self.make_temp_dir()
    sdk_location = 'http://storage.googleapis.com/my-gcs-bucket/tarball.tar.gz'

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).sdk_location = sdk_location

    def file_download(_, to_path):
      with open(to_path, 'w') as f:
        f.write('Package content.')
      return to_path

    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._download_file',
                    staticmethod(file_download)):
      self.assertEqual([
          names.STAGED_SDK_SOURCES_FILENAME,
          stager.SUBMISSION_ENV_DEPENDENCIES_FILE
      ],
                       self.stager.create_and_stage_job_resources(
                           options, staging_location=staging_dir)[1])

    tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME)
    with open(tarball_path) as f:
      self.assertEqual(f.read(), 'Package content.')

  def test_with_extra_packages(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()
    self.create_temp_file(os.path.join(source_dir, 'abc.tar.gz'), 'nothing')
    self.create_temp_file(os.path.join(source_dir, 'xyz.tar.gz'), 'nothing')
    self.create_temp_file(os.path.join(source_dir, 'xyz2.tar'), 'nothing')
    self.create_temp_file(os.path.join(source_dir, 'whl.whl'), 'nothing')
    self.create_temp_file(
        os.path.join(source_dir, stager.EXTRA_PACKAGES_FILE), 'nothing')

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).extra_packages = [
        os.path.join(source_dir, 'abc.tar.gz'),
        os.path.join(source_dir, 'xyz.tar.gz'),
        os.path.join(source_dir, 'xyz2.tar'),
        os.path.join(source_dir, 'whl.whl'),
        '/tmp/remote/remote_file.tar.gz'
    ]

    self.remote_copied_files = []

    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._download_file',
                    staticmethod(self.file_copy)):
      with mock.patch('apache_beam.runners.portability.stager_test'
                      '.stager.Stager._is_remote_path',
                      staticmethod(self.is_remote_path)):
        self.assertEqual([
            'abc.tar.gz',
            'xyz.tar.gz',
            'xyz2.tar',
            'whl.whl',
            'remote_file.tar.gz',
            stager.EXTRA_PACKAGES_FILE,
            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
        ],
                         self.stager.create_and_stage_job_resources(
                             options, staging_location=staging_dir)[1])
    with open(os.path.join(staging_dir, stager.EXTRA_PACKAGES_FILE)) as f:
      self.assertEqual([
          'abc.tar.gz\n',
          'xyz.tar.gz\n',
          'xyz2.tar\n',
          'whl.whl\n',
          'remote_file.tar.gz\n'
      ],
                       f.readlines())
    self.assertEqual(['/tmp/remote/remote_file.tar.gz'],
                     self.remote_copied_files)

  def test_with_extra_packages_missing_files(self):
    staging_dir = self.make_temp_dir()
    with self.assertRaises(RuntimeError) as cm:

      options = PipelineOptions()
      self.update_options(options)
      options.view_as(SetupOptions).extra_packages = ['nosuchfile.tar.gz']

      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertEqual(
        cm.exception.args[0],
        'The file %s cannot be found. It was specified in the '
        '--extra_packages command line option.' % 'nosuchfile.tar.gz')

  def test_with_extra_packages_invalid_file_name(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()
    self.create_temp_file(os.path.join(source_dir, 'abc.tgz'), 'nothing')
    with self.assertRaises(RuntimeError) as cm:
      options = PipelineOptions()
      self.update_options(options)
      options.view_as(SetupOptions).extra_packages = [
          os.path.join(source_dir, 'abc.tgz')
      ]
      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertEqual(
        cm.exception.args[0],
        'The --extra_package option expects a full path ending with '
        '".tar", ".tar.gz", ".whl" or ".zip" '
        'instead of %s' % os.path.join(source_dir, 'abc.tgz'))

  def test_with_jar_packages_missing_files(self):
    staging_dir = self.make_temp_dir()
    with self.assertRaises(RuntimeError) as cm:

      options = PipelineOptions()
      self.update_options(options)
      options.view_as(DebugOptions).experiments = [
          'jar_packages=nosuchfile.jar'
      ]
      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertEqual(
        cm.exception.args[0],
        'The file %s cannot be found. It was specified in the '
        '--experiment=\'jar_packages=\' command line option.' %
        'nosuchfile.jar')

  def test_with_jar_packages_invalid_file_name(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()
    self.create_temp_file(os.path.join(source_dir, 'abc.tgz'), 'nothing')
    with self.assertRaises(RuntimeError) as cm:
      options = PipelineOptions()
      self.update_options(options)
      options.view_as(DebugOptions).experiments = [
          'jar_packages=' + os.path.join(source_dir, 'abc.tgz')
      ]
      self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)
    self.assertEqual(
        cm.exception.args[0],
        'The --experiment=\'jar_packages=\' option expects a full path ending '
        'with ".jar" instead of %s' % os.path.join(source_dir, 'abc.tgz'))

  def test_with_jar_packages(self):
    staging_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()
    self.create_temp_file(os.path.join(source_dir, 'abc.jar'), 'nothing')
    self.create_temp_file(os.path.join(source_dir, 'xyz.jar'), 'nothing')
    self.create_temp_file(os.path.join(source_dir, 'ijk.jar'), 'nothing')

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(DebugOptions).experiments = [
        'jar_packages=%s,%s,%s,%s' % (
            os.path.join(source_dir, 'abc.jar'),
            os.path.join(source_dir, 'xyz.jar'),
            os.path.join(source_dir, 'ijk.jar'),
            '/tmp/remote/remote.jar')
    ]

    self.remote_copied_files = []

    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._download_file',
                    staticmethod(self.file_copy)):
      with mock.patch('apache_beam.runners.portability.stager_test'
                      '.stager.Stager._is_remote_path',
                      staticmethod(self.is_remote_path)):
        self.assertEqual([
            'abc.jar',
            'xyz.jar',
            'ijk.jar',
            'remote.jar',
            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
        ],
                         self.stager.create_and_stage_job_resources(
                             options, staging_location=staging_dir)[1])
    self.assertEqual(['/tmp/remote/remote.jar'], self.remote_copied_files)

  def test_remove_dependency_from_requirements(self):
    requirements_cache_dir = self.make_temp_dir()
    requirements = ['apache_beam\n', 'avro-python3\n', 'fastavro\n', 'numpy\n']
    with open(os.path.join(requirements_cache_dir, 'abc.txt'), 'w') as f:
      for i in range(len(requirements)):
        f.write(requirements[i])

    tmp_req_filename = self.stager._remove_dependency_from_requirements(
        requirements_file=os.path.join(requirements_cache_dir, 'abc.txt'),
        dependency_to_remove='apache_beam',
        temp_directory_path=requirements_cache_dir)
    with open(tmp_req_filename, 'r') as tf:
      lines = tf.readlines()
    self.assertEqual(['avro-python3\n', 'fastavro\n', 'numpy\n'], sorted(lines))

    tmp_req_filename = self.stager._remove_dependency_from_requirements(
        requirements_file=os.path.join(requirements_cache_dir, 'abc.txt'),
        dependency_to_remove='fastavro',
        temp_directory_path=requirements_cache_dir)

    with open(tmp_req_filename, 'r') as tf:
      lines = tf.readlines()
    self.assertEqual(['apache_beam\n', 'avro-python3\n', 'numpy\n'],
                     sorted(lines))

  def _populate_requitements_cache_fake(
      self, requirements_file, temp_dir, populate_cache_with_sdists):
    packages = set()
    if not populate_cache_with_sdists:
      self.create_temp_file(os.path.join(temp_dir, 'nothing.whl'), 'Fake whl')
      packages.add('nothing.whl')
    self.create_temp_file(
        os.path.join(temp_dir, 'nothing.tar.gz'), 'Fake tarball')
    packages.add('nothing.tar.gz')
    return packages

  # requirements cache will popultated with bdist/whl if present
  # else source would be downloaded.
  def test_populate_requirements_cache_with_bdist(self):
    staging_dir = self.make_temp_dir()
    requirements_cache_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)

    options.view_as(SetupOptions).requirements_cache = requirements_cache_dir
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
    # for default container image, the sdk_container_image option would be none
    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._populate_requirements_cache',
                    staticmethod(self._populate_requitements_cache_fake)):
      options.view_as(SetupOptions).requirements_cache_only_sources = False
      resources = self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)[1]
      for f in resources:
        if (f != stager.REQUIREMENTS_FILE and
            f != stager.SUBMISSION_ENV_DEPENDENCIES_FILE):
          self.assertTrue(('.tar.gz' in f) or ('.whl' in f))

  # requirements cache will populated only with sdists/sources
  def test_populate_requirements_cache_with_sdist(self):
    staging_dir = self.make_temp_dir()
    requirements_cache_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)

    options.view_as(SetupOptions).requirements_cache = requirements_cache_dir
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._populate_requirements_cache',
                    staticmethod(self._populate_requitements_cache_fake)):
      options.view_as(SetupOptions).requirements_cache_only_sources = True
      resources = self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)[1]

      for f in resources:
        if (f != stager.REQUIREMENTS_FILE and
            f != stager.SUBMISSION_ENV_DEPENDENCIES_FILE):
          self.assertTrue('.tar.gz' in f)
          self.assertTrue('.whl' not in f)

  def test_populate_requirements_cache_with_local_files(self):
    staging_dir = self.make_temp_dir()
    requirements_cache_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()
    pkg_dir = self.make_temp_dir()

    options = PipelineOptions()
    self.update_options(options)

    options.view_as(SetupOptions).requirements_cache = requirements_cache_dir
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    local_package = os.path.join(pkg_dir, 'local_package.tar.gz')
    self.create_temp_file(local_package, 'local-package-content')
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE),
        '\n'.join(['fake_pypi', local_package]))
    with mock.patch('apache_beam.runners.portability.stager_test'
                    '.stager.Stager._populate_requirements_cache',
                    staticmethod(self._populate_requitements_cache_fake)):
      options.view_as(SetupOptions).requirements_cache_only_sources = True
      resources = self.stager.create_and_stage_job_resources(
          options, staging_location=staging_dir)[1]

      self.assertEqual(
          sorted([
              stager.REQUIREMENTS_FILE,
              stager.EXTRA_PACKAGES_FILE,
              'nothing.tar.gz',
              'local_package.tar.gz',
              stager.SUBMISSION_ENV_DEPENDENCIES_FILE
          ]),
          sorted(resources))

      with open(os.path.join(staging_dir, stager.REQUIREMENTS_FILE)) as fin:
        requirements_contents = fin.read()
      self.assertIn('fake_pypi', requirements_contents)
      self.assertNotIn('local_package', requirements_contents)

      with open(os.path.join(staging_dir, stager.EXTRA_PACKAGES_FILE)) as fin:
        extra_packages_contents = fin.read()
      self.assertNotIn('fake_pypi', extra_packages_contents)
      self.assertIn('local_package', extra_packages_contents)

  def test_only_required_packages_staged_from_cache(self):
    """Test that only packages needed for current requirements are staged.

    This test verifies the fix for the issue where the entire cache directory
    was being staged, even when some packages weren't needed for the current
    workflow.
    """
    staging_dir = self.make_temp_dir()
    requirements_cache_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    # Pre-populate cache with packages from a "previous run"
    self.create_temp_file(
        os.path.join(requirements_cache_dir, 'old_package.whl'), 'old package')
    self.create_temp_file(
        os.path.join(requirements_cache_dir, 'another_old.tar.gz'), 'another')

    options = PipelineOptions()
    self.update_options(options)
    options.view_as(SetupOptions).requirements_cache = requirements_cache_dir
    options.view_as(SetupOptions).requirements_file = os.path.join(
        source_dir, stager.REQUIREMENTS_FILE)
    self.create_temp_file(
        os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'new_package')

    def populate_cache_for_new_requirements(
        requirements_file, cache_dir, populate_cache_with_sdists=False):
      # Simulate downloading only the packages needed for new requirements
      self.create_temp_file(
          os.path.join(cache_dir, 'new_package.whl'), 'new package content')
      return {'new_package.whl'}

    resources = self.stager.create_and_stage_job_resources(
        options,
        populate_requirements_cache=populate_cache_for_new_requirements,
        staging_location=staging_dir)[1]

    # Verify only new_package.whl is staged, not old packages
    self.assertIn('new_package.whl', resources)
    self.assertNotIn('old_package.whl', resources)
    self.assertNotIn('another_old.tar.gz', resources)

    # Verify the file exists in staging
    self.assertTrue(
        os.path.isfile(os.path.join(staging_dir, 'new_package.whl')))
    # Verify old packages are NOT in staging
    self.assertFalse(
        os.path.isfile(os.path.join(staging_dir, 'old_package.whl')))
    self.assertFalse(
        os.path.isfile(os.path.join(staging_dir, 'another_old.tar.gz')))

  def test_populate_requirements_cache_uses_find_links(self):
    """Test that _populate_requirements_cache uses --find-links to reuse cache.

    This test verifies that pip download is called with --find-links pointing
    to the cache directory, so packages already in cache are reused instead
    of being re-downloaded from PyPI.
    """
    requirements_cache_dir = self.make_temp_dir()
    source_dir = self.make_temp_dir()

    # Create a requirements file
    requirements_file = os.path.join(source_dir, 'requirements.txt')
    self.create_temp_file(requirements_file, 'some_package==1.0.0')

    captured_cmd_args = []

    def mock_check_output(cmd_args, **kwargs):
      captured_cmd_args.extend(cmd_args)
      return b''

    with mock.patch(
        'apache_beam.runners.portability.stager.processes.check_output',
        side_effect=mock_check_output):
      stager.Stager._populate_requirements_cache(
          requirements_file, requirements_cache_dir)

    # Verify --find-links is in the command with the cache directory
    self.assertIn('--find-links', captured_cmd_args)
    find_links_index = captured_cmd_args.index('--find-links')
    self.assertEqual(
        requirements_cache_dir, captured_cmd_args[find_links_index + 1])


class TestStager(stager.Stager):
  def stage_artifact(self, local_path_to_artifact, artifact_name, sha256):
    _LOGGER.info(
        'File copy from %s to %s.', local_path_to_artifact, artifact_name)
    shutil.copyfile(local_path_to_artifact, artifact_name)

  def commit_manifest(self):
    pass


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()
