# Path: fedjax/core/serialization.py
# def save_state(state, path):
# def load_state(path):
# def _ndarray_to_bytes(arr):
# def _dtype_from_name(name):
# def _ndarray_from_bytes(data):
# def _bytes_ndarray_to_bytes(x):
# def _object_ndarray_from_bytes(data):
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# def msgpack_serialize(pytree):
# def msgpack_deserialize(encoded_pytree):
# class _MsgpackExtType(enum.IntEnum):
# 
# Path: fedjax/core/sqlite_federated_data.py
# def decompress_and_deserialize(data: bytes):
#   def new(
#       path: str,
#       parse_examples: Callable[
#           [bytes], client_datasets.Examples] = decompress_and_deserialize
#   ) -> 'SQLiteFederatedData':
#   def __init__(self,
#                connection: sqlite3.Connection,
#                parse_examples: Callable[[bytes], client_datasets.Examples],
#                start: Optional[federated_data.ClientId] = None,
#                stop: Optional[federated_data.ClientId] = None,
#                preprocess_client: federated_data
#                .ClientPreprocessor = federated_data.NoOpClientPreprocessor,
#                preprocess_batch: client_datasets
#                .BatchPreprocessor = client_datasets.NoOpBatchPreprocessor):
#   def slice(
#       self,
#       start: Optional[federated_data.ClientId] = None,
#       stop: Optional[federated_data.ClientId] = None) -> 'SQLiteFederatedData':
#   def preprocess_client(
#       self, fn: Callable[[federated_data.ClientId, client_datasets.Examples],
#                          client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def preprocess_batch(
#       self, fn: Callable[[client_datasets.Examples], client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def _range_where(self) -> str:
#   def num_clients(self) -> int:
#   def client_ids(self) -> Iterator[federated_data.ClientId]:
#   def client_sizes(self) -> Iterator[Tuple[federated_data.ClientId, int]]:
#   def client_size(self, client_id: federated_data.ClientId) -> int:
#   def clients(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def _read_clients(self):
#   def shuffled_clients(
#       self,
#       buffer_size: int,
#       seed: Optional[int] = None
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_clients(
#       self, client_ids: Iterable[federated_data.ClientId]
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_client(
#       self,
#       client_id: federated_data.ClientId) -> client_datasets.ClientDataset:
#   def _client_dataset(self, client_id: federated_data.ClientId,
#                       data: bytes) -> client_datasets.ClientDataset:
#   def __init__(self, path: str):
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, exc_traceback):
#   def add_many(self,
#                client_ids_examples: Iterable[Tuple[bytes,
#                                                    client_datasets.Examples]]):
#     def prepare_parameters(ce):
#   def __init__(self, path: str,
#                parse_examples: Callable[[List[bytes]],
#                                         client_datasets.Examples],
#                split_name: str):
#   def __del__(self):
#   def __iter__(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def __next__(
#       self) -> Tuple[federated_data.ClientId, client_datasets.ClientDataset]:
# class SQLiteFederatedData(federated_data.FederatedData):
# class SQLiteFederatedDataBuilder(federated_data.FederatedDataBuilder):
# class TFFSQLiteClientsIterator:
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


import itertools
import os.path
import zlib
import numpy as np
import numpy.testing as npt
import sqlite3
from absl import flags
from absl.testing import absltest
from fedjax.core import serialization
from fedjax.core import sqlite_federated_data
from fedjax.core import util
# 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.
"""Tests for fedjax.core.sqlite_federated_data."""



tf = util.import_tf()

FLAGS = flags.FLAGS


class SQLiteFederatedDataTest(absltest.TestCase):

  @classmethod
  def setUpClass(cls):
    super().setUpClass()
    path = os.path.join(FLAGS.test_tmpdir, 'test_sqlite_federated_data.sqlite')

    # Create database file. First make sure the database file is empty.
    with open(path, 'w'):
      pass
    connection = sqlite3.connect(path)
    with connection:
      connection.execute("""
      CREATE TABLE federated_data (
        client_id BLOB NOT NULL PRIMARY KEY,
        data BLOB NOT NULL,

# Path: fedjax/core/serialization.py
# def save_state(state, path):
# def load_state(path):
# def _ndarray_to_bytes(arr):
# def _dtype_from_name(name):
# def _ndarray_from_bytes(data):
# def _bytes_ndarray_to_bytes(x):
# def _object_ndarray_from_bytes(data):
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# def msgpack_serialize(pytree):
# def msgpack_deserialize(encoded_pytree):
# class _MsgpackExtType(enum.IntEnum):
# 
# Path: fedjax/core/sqlite_federated_data.py
# def decompress_and_deserialize(data: bytes):
#   def new(
#       path: str,
#       parse_examples: Callable[
#           [bytes], client_datasets.Examples] = decompress_and_deserialize
#   ) -> 'SQLiteFederatedData':
#   def __init__(self,
#                connection: sqlite3.Connection,
#                parse_examples: Callable[[bytes], client_datasets.Examples],
#                start: Optional[federated_data.ClientId] = None,
#                stop: Optional[federated_data.ClientId] = None,
#                preprocess_client: federated_data
#                .ClientPreprocessor = federated_data.NoOpClientPreprocessor,
#                preprocess_batch: client_datasets
#                .BatchPreprocessor = client_datasets.NoOpBatchPreprocessor):
#   def slice(
#       self,
#       start: Optional[federated_data.ClientId] = None,
#       stop: Optional[federated_data.ClientId] = None) -> 'SQLiteFederatedData':
#   def preprocess_client(
#       self, fn: Callable[[federated_data.ClientId, client_datasets.Examples],
#                          client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def preprocess_batch(
#       self, fn: Callable[[client_datasets.Examples], client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def _range_where(self) -> str:
#   def num_clients(self) -> int:
#   def client_ids(self) -> Iterator[federated_data.ClientId]:
#   def client_sizes(self) -> Iterator[Tuple[federated_data.ClientId, int]]:
#   def client_size(self, client_id: federated_data.ClientId) -> int:
#   def clients(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def _read_clients(self):
#   def shuffled_clients(
#       self,
#       buffer_size: int,
#       seed: Optional[int] = None
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_clients(
#       self, client_ids: Iterable[federated_data.ClientId]
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_client(
#       self,
#       client_id: federated_data.ClientId) -> client_datasets.ClientDataset:
#   def _client_dataset(self, client_id: federated_data.ClientId,
#                       data: bytes) -> client_datasets.ClientDataset:
#   def __init__(self, path: str):
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, exc_traceback):
#   def add_many(self,
#                client_ids_examples: Iterable[Tuple[bytes,
#                                                    client_datasets.Examples]]):
#     def prepare_parameters(ce):
#   def __init__(self, path: str,
#                parse_examples: Callable[[List[bytes]],
#                                         client_datasets.Examples],
#                split_name: str):
#   def __del__(self):
#   def __iter__(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def __next__(
#       self) -> Tuple[federated_data.ClientId, client_datasets.ClientDataset]:
# class SQLiteFederatedData(federated_data.FederatedData):
# class SQLiteFederatedDataBuilder(federated_data.FederatedDataBuilder):
# class TFFSQLiteClientsIterator:
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


# Path: fedjax/core/sqlite_federated_data_test.py
import itertools
import os.path
import zlib
import numpy as np
import numpy.testing as npt
import sqlite3
from absl import flags
from absl.testing import absltest
from fedjax.core import serialization
from fedjax.core import sqlite_federated_data
from fedjax.core import util

# 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.
"""Tests for fedjax.core.sqlite_federated_data."""



tf = util.import_tf()

FLAGS = flags.FLAGS


class SQLiteFederatedDataTest(absltest.TestCase):

  @classmethod
  def setUpClass(cls):
    super().setUpClass()
    path = os.path.join(FLAGS.test_tmpdir, 'test_sqlite_federated_data.sqlite')

    # Create database file. First make sure the database file is empty.
    with open(path, 'w'):
      pass
    connection = sqlite3.connect(path)
    with connection:
      connection.execute("""
      CREATE TABLE federated_data (
        client_id BLOB NOT NULL PRIMARY KEY,
        data BLOB NOT NULL,

        num_examples INTEGER NOT NULL