# Path: fedjax/datasets/downloads.py
# def progress(n: int) -> Iterator[int]:
# def maybe_download(url: str,
#                    cache_dir: Optional[str] = None,
#                    progress_: Callable[[int], Iterator[int]] = progress) -> str:
# def format_duration(seconds: float) -> str:
# def log(*args, **kwargs):
# def default_cache_dir() -> str:
# def maybe_lzma_decompress(path) -> str:
# def validate_file(path: str, expected_num_bytes: int, expected_hexdigest: str):


import os
import os.path
import shutil
import sqlite3
from absl import flags
from absl.testing import absltest
from fedjax.datasets import downloads


class DownloadsTest(absltest.TestCase):
  SOURCE = 'https://storage.googleapis.com/tff-datasets-public/shakespeare.sqlite.lzma'
  HEXDIGEST = 'd3d11fceb9e105439ac6f4d52af6efafed5a2a1e1eb24c5bd2dd54ced242f5c4'
  NUM_BYTES = 1329828

  def setUp(self):
    super().setUp()
    # This will be set in actual tests.
    self._cache_dir = None

  def tearDown(self):
    super().tearDown()
    if self._cache_dir is not None:
      shutil.rmtree(self._cache_dir)

  def _with_cache_dir(self, cache_dir, expected_path):
    # First access, being copied.
    path = downloads.maybe_download(self.SOURCE, cache_dir)
    self._cache_dir = os.path.dirname(path)
    self.assertEqual(path, expected_path)
    downloads.validate_file(path, self.NUM_BYTES, self.HEXDIGEST)
    stat = os.stat(path)
    # ctime >= mtime because of renaming.
    self.assertGreaterEqual(stat.st_ctime_ns, stat.st_mtime_ns)

    # Second access, no copy.
    path1 = downloads.maybe_download(self.SOURCE, cache_dir)
    self.assertEqual(path1, path)

# Path: fedjax/datasets/downloads.py
# def progress(n: int) -> Iterator[int]:
# def maybe_download(url: str,
#                    cache_dir: Optional[str] = None,
#                    progress_: Callable[[int], Iterator[int]] = progress) -> str:
# def format_duration(seconds: float) -> str:
# def log(*args, **kwargs):
# def default_cache_dir() -> str:
# def maybe_lzma_decompress(path) -> str:
# def validate_file(path: str, expected_num_bytes: int, expected_hexdigest: str):


# Path: fedjax/datasets/downloads_test.py
import os
import os.path
import shutil
import sqlite3
from absl import flags
from absl.testing import absltest
from fedjax.datasets import downloads



class DownloadsTest(absltest.TestCase):
  SOURCE = 'https://storage.googleapis.com/tff-datasets-public/shakespeare.sqlite.lzma'
  HEXDIGEST = 'd3d11fceb9e105439ac6f4d52af6efafed5a2a1e1eb24c5bd2dd54ced242f5c4'
  NUM_BYTES = 1329828

  def setUp(self):
    super().setUp()
    # This will be set in actual tests.
    self._cache_dir = None

  def tearDown(self):
    super().tearDown()
    if self._cache_dir is not None:
      shutil.rmtree(self._cache_dir)

  def _with_cache_dir(self, cache_dir, expected_path):
    # First access, being copied.
    path = downloads.maybe_download(self.SOURCE, cache_dir)
    self._cache_dir = os.path.dirname(path)
    self.assertEqual(path, expected_path)
    downloads.validate_file(path, self.NUM_BYTES, self.HEXDIGEST)
    stat = os.stat(path)
    # ctime >= mtime because of renaming.
    self.assertGreaterEqual(stat.st_ctime_ns, stat.st_mtime_ns)

    # Second access, no copy.
    path1 = downloads.maybe_download(self.SOURCE, cache_dir)
    self.assertEqual(path1, path)

    downloads.validate_file(path1, self.NUM_BYTES, self.HEXDIGEST)