tests/test_hf_persistence.py

"""Unit tests for HFManager persistence logic."""

import os
import sys
import unittest
from unittest.mock import patch

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.hf_manager import HFManager


class TestHFPersistence(unittest.TestCase):
    def setUp(self):
        # Mock environment variable
        os.environ["HF_TOKEN"] = "fake_token"
        self.hf = HFManager()

    @patch("core.hf_manager.hf_hub_download")
    def test_download_file(self, mock_download):
        mock_download.return_value = "data/watchlist.json"

        result = self.hf.download_file("watchlist.json", target_dir="data")

        self.assertEqual(result, "data/watchlist.json")
        mock_download.assert_called_once()
        args, kwargs = mock_download.call_args
        self.assertEqual(kwargs["filename"], "watchlist.json")
        self.assertEqual(kwargs["repo_id"], self.hf.repo_id)

    @patch("core.hf_manager.HfApi.upload_file")
    @patch("os.path.exists")
    def test_upload_file(self, mock_exists, mock_upload):
        mock_exists.return_value = True

        success = self.hf.upload_file("data/watchlist.json", "watchlist.json")

        self.assertTrue(success)
        mock_upload.assert_called_once()
        args, kwargs = mock_upload.call_args
        self.assertEqual(kwargs["path_in_repo"], "watchlist.json")
        self.assertEqual(kwargs["repo_id"], self.hf.repo_id)

    def test_upload_file_no_token(self):
        self.hf.token = None
        success = self.hf.upload_file("data/watchlist.json")
        self.assertFalse(success)

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