tests/test_storage.py

from datetime import date
import json
import os
from pathlib import Path
from tempfile import TemporaryDirectory
import unittest
from unittest.mock import patch

from training_coach.models import CompletedSession, CompletedSet
from training_coach.storage import (
    HuggingFaceDatasetHistoryStore,
    JsonHistoryStore,
    create_history_store,
)


def _completed_session(day_number: int) -> CompletedSession:
    return CompletedSession(
        date=date(2026, 6, 11),
        day_number=day_number,
        completed_sets=[
            CompletedSet(
                exercise_id="dumbbell-row",
                set_number=1,
                actual_reps=16,
                actual_load=32.5,
                rpe=8,
            )
        ],
    )


class JsonHistoryStoreTest(unittest.TestCase):
    def test_missing_history_file_loads_as_empty_list(self):
        with TemporaryDirectory() as directory:
            store = JsonHistoryStore(Path(directory) / "history.json")

            self.assertEqual(store.load_completed_sessions(), [])

    def test_save_and_load_completed_sessions(self):
        with TemporaryDirectory() as directory:
            store = JsonHistoryStore(Path(directory) / "history.json")

            store.save_completed_sessions([_completed_session(day_number=2)])
            sessions = store.load_completed_sessions()

            self.assertEqual(len(sessions), 1)
            self.assertEqual(sessions[0].day_number, 2)
            self.assertEqual(sessions[0].completed_sets[0].exercise_id, "dumbbell-row")

    def test_append_completed_session_preserves_existing_history(self):
        with TemporaryDirectory() as directory:
            store = JsonHistoryStore(Path(directory) / "history.json")

            store.append_completed_session(_completed_session(day_number=1))
            store.append_completed_session(_completed_session(day_number=2))
            sessions = store.load_completed_sessions()

            self.assertEqual([session.day_number for session in sessions], [1, 2])


class _FakeHfApi:
    def __init__(self):
        self.created_repos = []
        self.uploads = []
        self.uploaded_payload = None

    def create_repo(self, **kwargs):
        self.created_repos.append(kwargs)

    def upload_file(self, **kwargs):
        self.uploads.append(kwargs)
        self.uploaded_payload = Path(kwargs["path_or_fileobj"]).read_text(
            encoding="utf-8"
        )


class HuggingFaceDatasetHistoryStoreTest(unittest.TestCase):
    def test_load_completed_sessions_downloads_history_file(self):
        with TemporaryDirectory() as directory:
            history_path = Path(directory) / "completed_sessions.json"
            history_path.write_text(
                json.dumps(
                    {
                        "completed_sessions": [
                            _completed_session(day_number=3).model_dump(mode="json")
                        ]
                    }
                ),
                encoding="utf-8",
            )
            store = HuggingFaceDatasetHistoryStore(
                repo_id="lucas/training-history",
                token="secret",
                downloader=lambda **_: str(history_path),
            )

            sessions = store.load_completed_sessions()

        self.assertEqual(len(sessions), 1)
        self.assertEqual(sessions[0].day_number, 3)

    def test_load_completed_sessions_forces_fresh_hub_download(self):
        download_calls = []

        with TemporaryDirectory() as directory:
            history_path = Path(directory) / "completed_sessions.json"
            history_path.write_text(
                json.dumps(
                    {
                        "completed_sessions": [
                            _completed_session(day_number=1).model_dump(mode="json")
                        ]
                    }
                ),
                encoding="utf-8",
            )

            def fake_downloader(**kwargs):
                download_calls.append(kwargs)
                return str(history_path)

            with patch("huggingface_hub.hf_hub_download", fake_downloader):
                store = HuggingFaceDatasetHistoryStore(
                    repo_id="lucas/training-history",
                    token="secret",
                    path_in_repo="history/completed_sessions.json",
                )
                store.load_completed_sessions()

        self.assertEqual(download_calls[0]["repo_id"], "lucas/training-history")
        self.assertEqual(download_calls[0]["filename"], "history/completed_sessions.json")
        self.assertEqual(download_calls[0]["repo_type"], "dataset")
        self.assertEqual(download_calls[0]["token"], "secret")
        self.assertEqual(download_calls[0]["force_download"], True)

    def test_save_completed_sessions_creates_private_dataset_and_uploads_file(self):
        api = _FakeHfApi()
        store = HuggingFaceDatasetHistoryStore(
            repo_id="lucas/training-history",
            token="secret",
            path_in_repo="history/completed_sessions.json",
            api=api,
        )

        store.save_completed_sessions([_completed_session(day_number=4)])

        self.assertEqual(api.created_repos[0]["repo_id"], "lucas/training-history")
        self.assertEqual(api.created_repos[0]["repo_type"], "dataset")
        self.assertEqual(api.created_repos[0]["private"], True)
        self.assertEqual(api.uploads[0]["path_in_repo"], "history/completed_sessions.json")
        self.assertEqual(api.uploads[0]["repo_type"], "dataset")
        payload = json.loads(api.uploaded_payload)
        self.assertEqual(payload["completed_sessions"][0]["day_number"], 4)


class HistoryStoreFactoryTest(unittest.TestCase):
    def test_create_history_store_defaults_to_local_json(self):
        with patch.dict(os.environ, {}, clear=True):
            store = create_history_store()

        self.assertIsInstance(store, JsonHistoryStore)

    def test_create_history_store_uses_hf_dataset_when_repo_is_set(self):
        with patch.dict(
            os.environ,
            {
                "HF_HISTORY_DATASET_REPO": "lucas/training-history",
                "HF_TOKEN": "secret",
                "HF_HISTORY_FILE": "history.json",
            },
            clear=True,
        ):
            store = create_history_store()

        self.assertIsInstance(store, HuggingFaceDatasetHistoryStore)
        self.assertEqual(store.repo_id, "lucas/training-history")
        self.assertEqual(store.token, "secret")
        self.assertEqual(store.path_in_repo, "history.json")


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