import pytest
from pydantic import BaseModel

import pys
from pys.pydantic import ModelWithID


@pytest.fixture
def storage():
    storage = pys.file_storage('storage.db')
    yield storage
    storage.destroy()


def test_sample(storage):
    # An author
    @pys.saveable
    class Author(BaseModel):
        name: str

    # And a book
    @pys.saveable
    class Book(BaseModel):
        title: str

    # A few books of Leo Tolstoy
    leo = Author(name='Leo Tolstoy')
    war_and_peace = Book(title='War and peace')
    for_kids = Book(title='For Kids')

    leo_id = storage.save(leo)
    storage.save(war_and_peace, (Author, leo_id))
    storage.save(for_kids, (Author, leo_id))

    leo_books = list(storage.list(Book, (Author, leo_id)))
    assert len(leo_books) == 2
    assert war_and_peace in leo_books
    assert for_kids in leo_books


def test_sample_model_with_id(storage):
    # An author
    class Author(ModelWithID):
        name: str

    # And a book
    class Book(ModelWithID):
        title: str

    # A few books of Leo Tolstoy
    leo = Author(name='Leo Tolstoy')
    war_and_peace = Book(title='War and peace')
    for_kids = Book(title='For Kids')

    storage.save(leo)
    storage.save(war_and_peace, leo)
    storage.save(for_kids, leo)

    leo_books = list(storage.list(Book, leo))
    assert len(leo_books) == 2
    assert war_and_peace in leo_books
    assert for_kids in leo_books
