dvc/tests/func/test_import_db.py

"""import_db tests without involving any databases."""

from agate import Table

data1 = Table([(i, i) for i in range(1, 10)], ["id", "value"])
data2 = Table([(i, i) for i in range(10, 20)], ["id", "value"])


def test_sql(mocker, tmp_dir, dvc):
    mocker.patch("dvc.utils.db.execute_sql", side_effect=[data1, data2])
    stage = dvc.imp_db(
        sql="select * from model", profile="profile", output_format="json"
    )

    db = {"file_format": "json", "profile": "profile", "query": "select * from model"}

    assert (tmp_dir / stage.relpath).parse() == {
        "md5": "dae475048115315c8c26a10538e4b2ec",
        "frozen": True,
        "deps": [{"db": db}],
        "outs": [
            {
                "hash": "md5",
                "md5": "56226d443067b195c90b427db557e1f2",
                "path": "results.json",
                "size": 243,
            }
        ],
    }

    dvc.update(stage.addressing)
    assert (tmp_dir / stage.relpath).parse() == {
        "md5": "79b42f9ddb2a3ed999eb494ccab470c6",
        "frozen": True,
        "deps": [{"db": db}],
        "outs": [
            {
                "hash": "md5",
                "md5": "08d806ee7e1f309c4561750805f02276",
                "path": "results.json",
                "size": 290,
            }
        ],
    }


def test_model(mocker, tmp_dir, dvc):
    mocker.patch("dvc.utils.db.get_model", side_effect=[data1, data2])
    stage = dvc.imp_db(model="model", output_format="json")

    assert (tmp_dir / stage.relpath).parse() == {
        "md5": "a8157fb5b28457e1910506e9d4d4fca3",
        "frozen": True,
        "deps": [{"db": {"file_format": "json", "model": "model"}}],
        "outs": [
            {
                "hash": "md5",
                "md5": "56226d443067b195c90b427db557e1f2",
                "path": "model.json",
                "size": 243,
            }
        ],
    }

    dvc.update(stage.addressing)
    assert (tmp_dir / stage.relpath).parse() == {
        "md5": "2f84287f109fc2b8755a9e64ad910b5d",
        "frozen": True,
        "deps": [{"db": {"file_format": "json", "model": "model"}}],
        "outs": [
            {
                "hash": "md5",
                "md5": "08d806ee7e1f309c4561750805f02276",
                "path": "model.json",
                "size": 290,
            }
        ],
    }