modelscope/tests/cli/test_download_cmd.py

import os
import os.path as osp
import shutil
import subprocess
import tempfile
import unittest
import uuid

from modelscope.hub.api import HubApi
from modelscope.hub.constants import Licenses, ModelVisibility
from modelscope.hub.repository import Repository
from modelscope.utils.logger import get_logger
from modelscope.utils.test_utils import (TEST_ACCESS_TOKEN1,
                                         TEST_MODEL_CHINESE_NAME,
                                         TEST_MODEL_ORG, test_level)

logger = get_logger()

DEFAULT_GIT_PATH = 'git'
download_model_file_name = 'test.bin'


class DownloadCMDTest(unittest.TestCase):

    def setUp(self):
        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
        self.tmp_dir = tempfile.TemporaryDirectory().name
        if not os.path.exists(self.tmp_dir):
            os.makedirs(self.tmp_dir)

        self.api = HubApi()
        self.api.login(TEST_ACCESS_TOKEN1)
        self.model_name = 'op-%s' % (uuid.uuid4().hex)
        self.model_id = '%s/%s' % (TEST_MODEL_ORG, self.model_name)
        self.revision = 'v0.1_test_revision'
        self.api.create_model(
            model_id=self.model_id,
            visibility=ModelVisibility.PUBLIC,
            license=Licenses.APACHE_V2,
            chinese_name=TEST_MODEL_CHINESE_NAME,
        )
        self.prepare_case()

    def prepare_case(self):
        temporary_dir = tempfile.mkdtemp()
        self.model_dir = os.path.join(temporary_dir, self.model_name)
        repo = Repository(self.model_dir, clone_from=self.model_id)
        os.system("echo 'testtest'>%s"
                  % os.path.join(self.model_dir, download_model_file_name))
        repo.push('add model')
        repo.tag_and_push(self.revision, 'Test revision')

    def tearDown(self):
        shutil.rmtree(self.tmp_dir)
        try:
            self.api.delete_model(model_id=self.model_id)
        except Exception as e:
            logger.warning(f'Error deleting model {self.model_id}: {e}')
        super().tearDown()

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download(self):
        cmd = f'python -m modelscope.cli.cli download --model {self.model_id}'
        stat, output = subprocess.getstatusoutput(cmd)
        self.assertEqual(stat, 0)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download_with_position_arg(self):
        cmd = f'python -m modelscope.cli.cli download {self.model_id}'
        stat, output = subprocess.getstatusoutput(cmd)
        self.assertEqual(stat, 0)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download_file(self):
        cmd = f'python -m modelscope.cli.cli download --model {self.model_id} {download_model_file_name}'
        stat, output = subprocess.getstatusoutput(cmd)
        self.assertEqual(stat, 0)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download_with_cache(self):
        cmd = f'python -m modelscope.cli.cli download --model {self.model_id} --cache_dir {self.tmp_dir}'
        stat, output = subprocess.getstatusoutput(cmd)
        if stat != 0:
            print(output)
        self.assertEqual(stat, 0)
        found = any(download_model_file_name in files
                    for _, _, files in os.walk(self.tmp_dir))
        self.assertTrue(
            found,
            f'{download_model_file_name} not found under {self.tmp_dir}')

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download_with_revision(self):
        cmd = f'python -m modelscope.cli.cli download --model {self.model_id} --revision {self.revision}'
        stat, output = subprocess.getstatusoutput(cmd)
        if stat != 0:
            print(output)
        self.assertEqual(stat, 0)


class DownloadCMDTokenTest(unittest.TestCase):

    def setUp(self):
        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
        self.api = HubApi()

        # Create private model repo
        self.model_id = '%s/%s' % (TEST_MODEL_ORG, 'test_model_with_token')
        self.api.create_repo(
            repo_id=self.model_id,
            repo_type='model',
            visibility='private',
            license=Licenses.APACHE_V2,
            chinese_name=TEST_MODEL_CHINESE_NAME,
            exist_ok=True,
        )

        # Create private dataset repo
        self.dataset_id = '%s/%s' % (TEST_MODEL_ORG, 'test_dataset_with_token')
        self.api.create_repo(
            repo_id=self.dataset_id,
            repo_type='dataset',
            visibility='private',
            license=Licenses.APACHE_V2,
            exist_ok=True,
        )

    def tearDown(self):
        try:
            self.api.delete_model(model_id=self.model_id)
        except Exception as e:
            logger.warning(f'Error deleting model {self.model_id}: {e}')
        try:
            self.api.delete_dataset(dataset_id=self.dataset_id)
        except Exception as e:
            logger.warning(f'Error deleting dataset {self.dataset_id}: {e}')
        super().tearDown()

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download_model_with_token(self):
        cmd = f'python -m modelscope.cli.cli download --model {self.model_id} --token {TEST_ACCESS_TOKEN1}'
        stat, output = subprocess.getstatusoutput(cmd)
        print(output)
        self.assertEqual(stat, 0)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_download_dataset_with_token(self):
        cmd = f'python -m modelscope.cli.cli download --dataset {self.dataset_id} --token {TEST_ACCESS_TOKEN1}'
        stat, output = subprocess.getstatusoutput(cmd)
        print(output)
        self.assertEqual(stat, 0)


if __name__ == '__main__':
    unittest.main()