import os
from click.testing import CliRunner
from tf_codage.cli import train_model

SAMPLE_NOTEBOOK = r"""
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# parameters\n",
    "N_EPOCHS = 1\n",
    "OUTPUT_DIR = ''\n",
    "LEARNING_RATE = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "with open(os.path.join(OUTPUT_DIR, 'test.txt'), 'w') as fid:\n",
    "    fid.write(f\"LEARNING_RATE={LEARNING_RATE}\")\n",
    "    fid.write(f\"N_EPOCHS={N_EPOCHS}\")\n",
    "    fid.write(f\"VERBOSITY={VERBOSITY}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
"""


def test_train_model_help():
    """test help message of train_model command"""
    runner = CliRunner()
    result = runner.invoke(train_model.main, ["--help"])
    assert result.exit_code == 0
    assert result.output.startswith("Usage:")
    assert "NOTEBOOK" in result.output


def test_train_model_output():
    """test files generated by train_model from input notebook"""
    runner = CliRunner()
    with runner.isolated_filesystem():
        with open("hello.ipynb", "w") as f:
            f.write(SAMPLE_NOTEBOOK)

        result = runner.invoke(
            train_model.main,
            [
                "hello.ipynb",
                "-n",
                "0",
                "--save-path",
                "models",
                "--log-path",
                "logs",
                "-l",
                "0.123",
                "--gpu",
                "none",
            ],
        )
        assert result.exit_code == 0, result.output
        with open(
            os.path.join("models", "hello-n_epochs-0-learning_rate-0.123", "test.txt")
        ) as fid:
            model_output = fid.read()
        assert "LEARNING_RATE=0.123" in model_output
        assert "N_EPOCHS=0" in model_output
        assert "VERBOSITY=2" in model_output
