anki-sync-server/tests/collection_test_base.py

# -*- coding: utf-8 -*-
import unittest
import tempfile
import os
from unittest.mock import MagicMock
import shutil

import anki
import anki.storage

from ankisyncd.collection import CollectionManager


class CollectionTestBase(unittest.TestCase):
    """Parent class for tests that need a collection set up and torn down."""

    def setUp(self):
        self.temp_dir = tempfile.mkdtemp()
        self.collection_path = os.path.join(self.temp_dir, "collection.anki2")
        cm = CollectionManager({})
        collectionWrapper = cm.get_collection(self.collection_path)
        self.collection = collectionWrapper._get_collection()
        self.mock_app = MagicMock()

    def tearDown(self):
        self.collection.close()
        self.collection = None
        shutil.rmtree(self.temp_dir)
        self.mock_app.reset_mock()

    # TODO: refactor into some kind of utility
    def add_note(self, data):
        from anki.notes import Note

        model = self.collection.models.byName(data["model"])

        note = Note(self.collection, model)
        for name, value in data["fields"].items():
            note[name] = value

        if "tags" in data:
            note.setTagsFromStr(data["tags"])

        self.collection.addNote(note)

    # TODO: refactor into a parent class
    def add_default_note(self, count=1):
        data = {
            "model": "Basic",
            "fields": {
                "Front": "The front",
                "Back": "The back",
            },
            "tags": "Tag1 Tag2",
        }
        for idx in range(0, count):
            self.add_note(data)