From 7983f74b2e417a2cdf81318fcd3d0b6d207e0406 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 15:39:05 +0630
Subject: [PATCH 01/12] test: add user tests

---
 .../template-common/requirements.txt          |  4 ++
 .../requirements.txt                          |  4 ++
 .../test_all.py                               | 70 ++++++++++++++++++
 .../template-vision-dcgan/requirements.txt    |  4 ++
 .../template-vision-dcgan/test_all.py         | 72 +++++++++++++++++++
 .../requirements.txt                          |  4 ++
 .../template-vision-segmentation/test_all.py  | 70 ++++++++++++++++++
 7 files changed, 228 insertions(+)
 create mode 100644 src/templates/template-vision-classification/test_all.py
 create mode 100644 src/templates/template-vision-dcgan/test_all.py
 create mode 100644 src/templates/template-vision-segmentation/test_all.py

diff --git a/src/templates/template-common/requirements.txt b/src/templates/template-common/requirements.txt
index 4bbe4290..52dc19aa 100644
--- a/src/templates/template-common/requirements.txt
+++ b/src/templates/template-common/requirements.txt
@@ -12,3 +12,7 @@ pyyaml
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-classification/requirements.txt b/src/templates/template-vision-classification/requirements.txt
index 4bbe4290..52dc19aa 100644
--- a/src/templates/template-vision-classification/requirements.txt
+++ b/src/templates/template-vision-classification/requirements.txt
@@ -12,3 +12,7 @@ pyyaml
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-classification/test_all.py b/src/templates/template-vision-classification/test_all.py
new file mode 100644
index 00000000..a130e1c7
--- /dev/null
+++ b/src/templates/template-vision-classification/test_all.py
@@ -0,0 +1,70 @@
+import os
+from argparse import Namespace
+from numbers import Number
+from typing import Iterable
+
+import ignite.distributed as idist
+import pytest
+import torch
+from datasets import get_datasets
+from ignite.contrib.handlers.param_scheduler import ParamScheduler
+from ignite.engine import Engine
+from torch import nn, optim
+from torch.functional import Tensor
+from torch.optim.lr_scheduler import _LRScheduler
+from torch.utils.data import Dataset
+from trainers import evaluate_function
+from utils import initialize
+
+
+def set_up():
+    model = nn.Linear(1, 1)
+    optimizer = optim.Adam(model.parameters())
+    device = idist.device()
+    loss_fn = nn.MSELoss()
+    batch = [torch.tensor([1.0]), torch.tensor([1.0])]
+
+    return model, optimizer, device, loss_fn, batch
+
+
+@pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
+def test_get_datasets(tmp_path):
+    train_ds, eval_ds = get_datasets(tmp_path)
+
+    assert isinstance(train_ds, Dataset)
+    assert isinstance(eval_ds, Dataset)
+    train_batch = next(iter(train_ds))
+    assert isinstance(train_batch, Iterable)
+    assert isinstance(train_batch[0], Tensor)
+    assert isinstance(train_batch[1], Number)
+    assert train_batch[0].ndim == 3
+    eval_batch = next(iter(eval_ds))
+    assert isinstance(eval_batch, Iterable)
+    assert isinstance(eval_batch[0], Tensor)
+    assert isinstance(eval_batch[1], Number)
+    assert eval_batch[0].ndim == 3
+
+
+def test_evaluate_fn():
+    model, _, device, _, batch = set_up()
+    engine = Engine(lambda e, b: 1)
+    config = Namespace(use_amp=False)
+    output = evaluate_function(config, engine, batch, model, device)
+    assert isinstance(output, tuple)
+
+
+def test_initialize():
+    config = Namespace(
+        model="squeezenet1_0",
+        lr=1e-3,
+        momentum=0.9,
+        weight_decay=1e-4,
+        num_iters_per_epoch=1,
+        num_warmup_epochs=1,
+        max_epochs=1,
+    )
+    model, optimizer, loss_fn, lr_scheduler = initialize(config)
+    assert isinstance(model, nn.Module)
+    assert isinstance(optimizer, optim.Optimizer)
+    assert isinstance(loss_fn, nn.Module)
+    assert isinstance(lr_scheduler, (_LRScheduler, ParamScheduler))
diff --git a/src/templates/template-vision-dcgan/requirements.txt b/src/templates/template-vision-dcgan/requirements.txt
index 4bbe4290..52dc19aa 100644
--- a/src/templates/template-vision-dcgan/requirements.txt
+++ b/src/templates/template-vision-dcgan/requirements.txt
@@ -12,3 +12,7 @@ pyyaml
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
new file mode 100644
index 00000000..a65ecf98
--- /dev/null
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -0,0 +1,72 @@
+import os
+from argparse import Namespace
+from numbers import Number
+from typing import Iterable
+
+import ignite.distributed as idist
+import pytest
+import torch
+from datasets import get_datasets
+from ignite.engine import Engine
+from models import Discriminator, Generator
+from torch import nn, optim
+from torch.functional import Tensor
+from torch.utils.data import Dataset
+from trainers import train_function
+
+
+def set_up():
+    model = nn.Linear(1, 1)
+    optimizer = optim.Adam(model.parameters())
+    device = idist.device()
+    loss_fn = nn.MSELoss()
+    batch = [torch.tensor([1.0]), torch.tensor([1.0])]
+
+    return model, optimizer, device, loss_fn, batch
+
+
+@pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
+def test_get_datasets(tmp_path):
+    dataset, _ = get_datasets("cifar10", tmp_path)
+
+    assert isinstance(dataset, Dataset)
+    batch = next(iter(dataset))
+    assert isinstance(batch, Iterable)
+    assert isinstance(batch[0], Tensor)
+    assert isinstance(batch[1], Number)
+    assert batch[0].ndim == 3
+
+
+def test_models():
+    model_G = Generator(100, 64, 3)
+    model_D = Discriminator(3, 64)
+    x = torch.rand([1, 100, 32, 32])
+    y = model_G(x)
+    y.sum().backward()
+    z = model_D(y)
+    assert y.shape == torch.Size([1, 3, 560, 560])
+    assert z.shape == torch.Size([1024])
+    assert isinstance(model_D, nn.Module)
+    assert isinstance(model_G, nn.Module)
+
+
+def test_train_fn():
+    model, optimizer, device, loss_fn, batch = set_up()
+    real_labels = torch.ones(2, device=device)
+    fake_labels = torch.zeros(2, device=device)
+    engine = Engine(lambda e, b: 1)
+    config = Namespace(use_amp=False, batch_size=2, z_dim=100)
+    output = train_function(
+        config,
+        engine,
+        batch,
+        model,
+        model,
+        loss_fn,
+        optimizer,
+        optimizer,
+        device,
+        real_labels,
+        fake_labels,
+    )
+    assert isinstance(output, dict)
diff --git a/src/templates/template-vision-segmentation/requirements.txt b/src/templates/template-vision-segmentation/requirements.txt
index 22d3c3d5..054bac44 100644
--- a/src/templates/template-vision-segmentation/requirements.txt
+++ b/src/templates/template-vision-segmentation/requirements.txt
@@ -14,3 +14,7 @@ image_dataset_viz
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
new file mode 100644
index 00000000..a130e1c7
--- /dev/null
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -0,0 +1,70 @@
+import os
+from argparse import Namespace
+from numbers import Number
+from typing import Iterable
+
+import ignite.distributed as idist
+import pytest
+import torch
+from datasets import get_datasets
+from ignite.contrib.handlers.param_scheduler import ParamScheduler
+from ignite.engine import Engine
+from torch import nn, optim
+from torch.functional import Tensor
+from torch.optim.lr_scheduler import _LRScheduler
+from torch.utils.data import Dataset
+from trainers import evaluate_function
+from utils import initialize
+
+
+def set_up():
+    model = nn.Linear(1, 1)
+    optimizer = optim.Adam(model.parameters())
+    device = idist.device()
+    loss_fn = nn.MSELoss()
+    batch = [torch.tensor([1.0]), torch.tensor([1.0])]
+
+    return model, optimizer, device, loss_fn, batch
+
+
+@pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
+def test_get_datasets(tmp_path):
+    train_ds, eval_ds = get_datasets(tmp_path)
+
+    assert isinstance(train_ds, Dataset)
+    assert isinstance(eval_ds, Dataset)
+    train_batch = next(iter(train_ds))
+    assert isinstance(train_batch, Iterable)
+    assert isinstance(train_batch[0], Tensor)
+    assert isinstance(train_batch[1], Number)
+    assert train_batch[0].ndim == 3
+    eval_batch = next(iter(eval_ds))
+    assert isinstance(eval_batch, Iterable)
+    assert isinstance(eval_batch[0], Tensor)
+    assert isinstance(eval_batch[1], Number)
+    assert eval_batch[0].ndim == 3
+
+
+def test_evaluate_fn():
+    model, _, device, _, batch = set_up()
+    engine = Engine(lambda e, b: 1)
+    config = Namespace(use_amp=False)
+    output = evaluate_function(config, engine, batch, model, device)
+    assert isinstance(output, tuple)
+
+
+def test_initialize():
+    config = Namespace(
+        model="squeezenet1_0",
+        lr=1e-3,
+        momentum=0.9,
+        weight_decay=1e-4,
+        num_iters_per_epoch=1,
+        num_warmup_epochs=1,
+        max_epochs=1,
+    )
+    model, optimizer, loss_fn, lr_scheduler = initialize(config)
+    assert isinstance(model, nn.Module)
+    assert isinstance(optimizer, optim.Optimizer)
+    assert isinstance(loss_fn, nn.Module)
+    assert isinstance(lr_scheduler, (_LRScheduler, ParamScheduler))

From 2acd7939429cb495d8d93205d11093f15cb98524 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 16:16:11 +0630
Subject: [PATCH 02/12] test: refactor tests, add `include_test` checkbox

- add CI
- add checks in generating zip tests

ref #103
---
 .github/workflows/ci.yml                      | 10 +--
 __tests__/vision-classification.spec.js       |  3 +
 __tests__/vision-dcgan.spec.js                |  3 +
 __tests__/vision-segmentation.spec.js         |  3 +
 scripts/requirements.txt                      |  1 +
 scripts/run_tests.sh                          |  1 +
 src/components/TabTemplates.vue               |  9 ++-
 src/store.js                                  |  6 +-
 .../test_all.py                               | 61 +++++++-----------
 .../template-vision-dcgan/test_all.py         | 63 +++++++++----------
 .../template-vision-segmentation/test_all.py  | 61 +++++++-----------
 src/templates/templates.json                  |  3 +
 12 files changed, 107 insertions(+), 117 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 7422e483..9da60704 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -10,6 +10,8 @@ on:
 jobs:
   tests:
     runs-on: ${{ matrix.os }}
+    env:
+      RUN_SLOW_TESTS: 1
     strategy:
       fail-fast: false
       matrix:
@@ -59,14 +61,14 @@ jobs:
       - run: pnpm test:ci
       - run: sh ./scripts/run_tests.sh unzip
 
-      - name: Run simple
-        run: sh ./scripts/run_tests.sh simple
-        if: github.event_name != 'schedule'
-
       - name: Run all
         run: sh ./scripts/run_tests.sh all
         if: github.event_name != 'schedule'
 
+      - name: Run simple
+        run: sh ./scripts/run_tests.sh simple
+        if: github.event_name != 'schedule'
+
       - name: Run launch
         run: sh ./scripts/run_tests.sh launch
         if: github.event_name == 'schedule'
diff --git a/__tests__/vision-classification.spec.js b/__tests__/vision-classification.spec.js
index 13cdfc73..c2c59b2b 100644
--- a/__tests__/vision-classification.spec.js
+++ b/__tests__/vision-classification.spec.js
@@ -43,6 +43,9 @@ test('vision classification simple', async () => {
 test('vision classification all', async () => {
   await page.selectOption('select', 'template-vision-classification')
 
+  await page.check('#include_test-checkbox')
+  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
   await page.waitForSelector('text=README.md')
   await page.click(':nth-match(:text("Training"), 2)')
 
diff --git a/__tests__/vision-dcgan.spec.js b/__tests__/vision-dcgan.spec.js
index 81753a2e..127458f2 100644
--- a/__tests__/vision-dcgan.spec.js
+++ b/__tests__/vision-dcgan.spec.js
@@ -43,6 +43,9 @@ test('vision dcgan simple', async () => {
 test('vision dcgan all', async () => {
   await page.selectOption('select', 'template-vision-dcgan')
 
+  await page.check('#include_test-checkbox')
+  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
   await page.waitForSelector('text=README.md')
   await page.click(':nth-match(:text("Training"), 2)')
 
diff --git a/__tests__/vision-segmentation.spec.js b/__tests__/vision-segmentation.spec.js
index 84ab3e10..ef7c2fff 100644
--- a/__tests__/vision-segmentation.spec.js
+++ b/__tests__/vision-segmentation.spec.js
@@ -43,6 +43,9 @@ test('vision segmentation simple', async () => {
 test('vision segmentation all', async () => {
   await page.selectOption('select', 'template-vision-segmentation')
 
+  await page.check('#include_test-checkbox')
+  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
   await page.waitForSelector('text=README.md')
   await page.click(':nth-match(:text("Training"), 2)')
 
diff --git a/scripts/requirements.txt b/scripts/requirements.txt
index 7a659699..e0c29e52 100644
--- a/scripts/requirements.txt
+++ b/scripts/requirements.txt
@@ -4,3 +4,4 @@ pytorch-ignite>=0.4.2
 pyyaml
 albumentations
 image_dataset_viz
+pytest
diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh
index 84eeee1e..19385ad7 100644
--- a/scripts/run_tests.sh
+++ b/scripts/run_tests.sh
@@ -27,6 +27,7 @@ run_all() {
   for dir in $(find ./dist-tests/*-all -type d)
   do
     cd $dir
+    pytest -vra --color=yes --tb=short test_*.py
     python main.py --data_path ~/data
     cd $CWD
   done
diff --git a/src/components/TabTemplates.vue b/src/components/TabTemplates.vue
index c6f298f8..d0847c83 100644
--- a/src/components/TabTemplates.vue
+++ b/src/components/TabTemplates.vue
@@ -10,16 +10,23 @@
         @change.prevent="downloadTemplates"
       />
     </div>
+    <div class="include-test">
+      <FormCheckbox
+        label="Include a test file for the generated code"
+        saveKey="include_test"
+      />
+    </div>
   </div>
 </template>
 
 <script>
 import FormSelect from './FormSelect.vue'
+import FormCheckbox from './FormCheckbox.vue'
 import templates from '../templates/templates.json'
 import { store, fetchTemplates } from '../store.js'
 
 export default {
-  components: { FormSelect },
+  components: { FormSelect, FormCheckbox },
   setup() {
     const templateLabel = 'Choose A Template'
     const templateOptions = Object.keys(templates)
diff --git a/src/store.js b/src/store.js
index 106660b3..855afe14 100644
--- a/src/store.js
+++ b/src/store.js
@@ -40,7 +40,8 @@ export const msg = reactive({
 export const store = reactive({
   code: {},
   config: {
-    template: ''
+    template: '',
+    include_test: false
   }
 })
 
@@ -61,6 +62,9 @@ export function genCode() {
   const currentFiles = files[store.config.template]
   if (currentFiles && Object.keys(currentFiles).length) {
     for (const file in currentFiles) {
+      if (!store.config.include_test && file === 'test_all.py') {
+        continue
+      }
       store.code[file] = ejs
         .render(currentFiles[file], store.config)
         .replaceAll(/(\n\n\n\n)+/gi, '\n')
diff --git a/src/templates/template-vision-classification/test_all.py b/src/templates/template-vision-classification/test_all.py
index a130e1c7..b3acc52d 100644
--- a/src/templates/template-vision-classification/test_all.py
+++ b/src/templates/template-vision-classification/test_all.py
@@ -6,15 +6,10 @@
 import ignite.distributed as idist
 import pytest
 import torch
-from datasets import get_datasets
-from ignite.contrib.handlers.param_scheduler import ParamScheduler
-from ignite.engine import Engine
-from torch import nn, optim
-from torch.functional import Tensor
-from torch.optim.lr_scheduler import _LRScheduler
-from torch.utils.data import Dataset
-from trainers import evaluate_function
-from utils import initialize
+from data import setup_data
+from torch import Tensor, nn, optim
+from torch.utils.data.dataloader import DataLoader
+from trainers import setup_evaluator
 
 
 def set_up():
@@ -27,44 +22,32 @@ def set_up():
     return model, optimizer, device, loss_fn, batch
 
 
-@pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
-def test_get_datasets(tmp_path):
-    train_ds, eval_ds = get_datasets(tmp_path)
+@pytest.mark.skipif(
+    os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests"
+)
+def test_setup_data():
+    config = Namespace(
+        data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
+    )
+    dataloader_train, dataloader_eval = setup_data(config)
 
-    assert isinstance(train_ds, Dataset)
-    assert isinstance(eval_ds, Dataset)
-    train_batch = next(iter(train_ds))
+    assert isinstance(dataloader_train, DataLoader)
+    assert isinstance(dataloader_eval, DataLoader)
+    train_batch = next(dataloader_train)
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
     assert isinstance(train_batch[1], Number)
-    assert train_batch[0].ndim == 3
-    eval_batch = next(iter(eval_ds))
+    assert train_batch[0].ndim == 4
+    eval_batch = next(dataloader_eval)
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
     assert isinstance(eval_batch[1], Number)
-    assert eval_batch[0].ndim == 3
+    assert eval_batch[0].ndim == 4
 
 
-def test_evaluate_fn():
+def test_setup_evaluator():
     model, _, device, _, batch = set_up()
-    engine = Engine(lambda e, b: 1)
     config = Namespace(use_amp=False)
-    output = evaluate_function(config, engine, batch, model, device)
-    assert isinstance(output, tuple)
-
-
-def test_initialize():
-    config = Namespace(
-        model="squeezenet1_0",
-        lr=1e-3,
-        momentum=0.9,
-        weight_decay=1e-4,
-        num_iters_per_epoch=1,
-        num_warmup_epochs=1,
-        max_epochs=1,
-    )
-    model, optimizer, loss_fn, lr_scheduler = initialize(config)
-    assert isinstance(model, nn.Module)
-    assert isinstance(optimizer, optim.Optimizer)
-    assert isinstance(loss_fn, nn.Module)
-    assert isinstance(lr_scheduler, (_LRScheduler, ParamScheduler))
+    evaluator = setup_evaluator(config, model, device)
+    evaluator.run(batch)
+    assert isinstance(evaluator.state.output, tuple)
diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
index a65ecf98..27d8bc05 100644
--- a/src/templates/template-vision-dcgan/test_all.py
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -6,13 +6,11 @@
 import ignite.distributed as idist
 import pytest
 import torch
-from datasets import get_datasets
-from ignite.engine import Engine
+from data import setup_data
 from models import Discriminator, Generator
-from torch import nn, optim
-from torch.functional import Tensor
-from torch.utils.data import Dataset
-from trainers import train_function
+from torch import Tensor, nn, optim
+from torch.utils.data.dataloader import DataLoader
+from trainers import setup_trainer
 
 
 def set_up():
@@ -25,16 +23,27 @@ def set_up():
     return model, optimizer, device, loss_fn, batch
 
 
-@pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
-def test_get_datasets(tmp_path):
-    dataset, _ = get_datasets("cifar10", tmp_path)
+@pytest.mark.skipif(
+    os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests"
+)
+def test_setup_data():
+    config = Namespace(
+        data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
+    )
+    dataloader_train, dataloader_eval = setup_data(config)
 
-    assert isinstance(dataset, Dataset)
-    batch = next(iter(dataset))
-    assert isinstance(batch, Iterable)
-    assert isinstance(batch[0], Tensor)
-    assert isinstance(batch[1], Number)
-    assert batch[0].ndim == 3
+    assert isinstance(dataloader_train, DataLoader)
+    assert isinstance(dataloader_eval, DataLoader)
+    train_batch = next(dataloader_train)
+    assert isinstance(train_batch, Iterable)
+    assert isinstance(train_batch[0], Tensor)
+    assert isinstance(train_batch[1], Number)
+    assert train_batch[0].ndim == 4
+    eval_batch = next(dataloader_eval)
+    assert isinstance(eval_batch, Iterable)
+    assert isinstance(eval_batch[0], Tensor)
+    assert isinstance(eval_batch[1], Number)
+    assert eval_batch[0].ndim == 4
 
 
 def test_models():
@@ -50,23 +59,11 @@ def test_models():
     assert isinstance(model_G, nn.Module)
 
 
-def test_train_fn():
+def test_setup_trainer():
     model, optimizer, device, loss_fn, batch = set_up()
-    real_labels = torch.ones(2, device=device)
-    fake_labels = torch.zeros(2, device=device)
-    engine = Engine(lambda e, b: 1)
-    config = Namespace(use_amp=False, batch_size=2, z_dim=100)
-    output = train_function(
-        config,
-        engine,
-        batch,
-        model,
-        model,
-        loss_fn,
-        optimizer,
-        optimizer,
-        device,
-        real_labels,
-        fake_labels,
+    config = Namespace(use_amp=False, train_batch_size=2, z_dim=100)
+    trainer = setup_trainer(
+        config, model, model, optimizer, optimizer, loss_fn, device
     )
-    assert isinstance(output, dict)
+    trainer.run(batch)
+    assert isinstance(trainer.state.output, dict)
diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index a130e1c7..b3acc52d 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -6,15 +6,10 @@
 import ignite.distributed as idist
 import pytest
 import torch
-from datasets import get_datasets
-from ignite.contrib.handlers.param_scheduler import ParamScheduler
-from ignite.engine import Engine
-from torch import nn, optim
-from torch.functional import Tensor
-from torch.optim.lr_scheduler import _LRScheduler
-from torch.utils.data import Dataset
-from trainers import evaluate_function
-from utils import initialize
+from data import setup_data
+from torch import Tensor, nn, optim
+from torch.utils.data.dataloader import DataLoader
+from trainers import setup_evaluator
 
 
 def set_up():
@@ -27,44 +22,32 @@ def set_up():
     return model, optimizer, device, loss_fn, batch
 
 
-@pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
-def test_get_datasets(tmp_path):
-    train_ds, eval_ds = get_datasets(tmp_path)
+@pytest.mark.skipif(
+    os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests"
+)
+def test_setup_data():
+    config = Namespace(
+        data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
+    )
+    dataloader_train, dataloader_eval = setup_data(config)
 
-    assert isinstance(train_ds, Dataset)
-    assert isinstance(eval_ds, Dataset)
-    train_batch = next(iter(train_ds))
+    assert isinstance(dataloader_train, DataLoader)
+    assert isinstance(dataloader_eval, DataLoader)
+    train_batch = next(dataloader_train)
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
     assert isinstance(train_batch[1], Number)
-    assert train_batch[0].ndim == 3
-    eval_batch = next(iter(eval_ds))
+    assert train_batch[0].ndim == 4
+    eval_batch = next(dataloader_eval)
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
     assert isinstance(eval_batch[1], Number)
-    assert eval_batch[0].ndim == 3
+    assert eval_batch[0].ndim == 4
 
 
-def test_evaluate_fn():
+def test_setup_evaluator():
     model, _, device, _, batch = set_up()
-    engine = Engine(lambda e, b: 1)
     config = Namespace(use_amp=False)
-    output = evaluate_function(config, engine, batch, model, device)
-    assert isinstance(output, tuple)
-
-
-def test_initialize():
-    config = Namespace(
-        model="squeezenet1_0",
-        lr=1e-3,
-        momentum=0.9,
-        weight_decay=1e-4,
-        num_iters_per_epoch=1,
-        num_warmup_epochs=1,
-        max_epochs=1,
-    )
-    model, optimizer, loss_fn, lr_scheduler = initialize(config)
-    assert isinstance(model, nn.Module)
-    assert isinstance(optimizer, optim.Optimizer)
-    assert isinstance(loss_fn, nn.Module)
-    assert isinstance(lr_scheduler, (_LRScheduler, ParamScheduler))
+    evaluator = setup_evaluator(config, model, device)
+    evaluator.run(batch)
+    assert isinstance(evaluator.state.output, tuple)
diff --git a/src/templates/templates.json b/src/templates/templates.json
index d1d57867..248b029d 100644
--- a/src/templates/templates.json
+++ b/src/templates/templates.json
@@ -7,6 +7,7 @@
     "model.py",
     "trainers.py",
     "utils.py",
+    "test_all.py",
     "requirements.txt"
   ],
   "template-vision-dcgan": [
@@ -17,6 +18,7 @@
     "model.py",
     "trainers.py",
     "utils.py",
+    "test_all.py",
     "requirements.txt"
   ],
   "template-vision-segmentation": [
@@ -27,6 +29,7 @@
     "trainers.py",
     "utils.py",
     "vis.py",
+    "test_all.py",
     "requirements.txt"
   ]
 }

From 1cde80873b7a275ed30198ed191a8277994e23f9 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 16:31:20 +0630
Subject: [PATCH 03/12] fix: test

---
 src/templates/template-vision-classification/test_all.py | 6 +++---
 src/templates/template-vision-dcgan/test_all.py          | 6 +++---
 src/templates/template-vision-segmentation/test_all.py   | 6 +++---
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/templates/template-vision-classification/test_all.py b/src/templates/template-vision-classification/test_all.py
index b3acc52d..abf2b927 100644
--- a/src/templates/template-vision-classification/test_all.py
+++ b/src/templates/template-vision-classification/test_all.py
@@ -33,12 +33,12 @@ def test_setup_data():
 
     assert isinstance(dataloader_train, DataLoader)
     assert isinstance(dataloader_eval, DataLoader)
-    train_batch = next(dataloader_train)
+    train_batch = next(iter(dataloader_train))
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
     assert isinstance(train_batch[1], Number)
     assert train_batch[0].ndim == 4
-    eval_batch = next(dataloader_eval)
+    eval_batch = next(iter(dataloader_eval))
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
     assert isinstance(eval_batch[1], Number)
@@ -49,5 +49,5 @@ def test_setup_evaluator():
     model, _, device, _, batch = set_up()
     config = Namespace(use_amp=False)
     evaluator = setup_evaluator(config, model, device)
-    evaluator.run(batch)
+    evaluator.run([batch, batch])
     assert isinstance(evaluator.state.output, tuple)
diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
index 27d8bc05..09a1cac8 100644
--- a/src/templates/template-vision-dcgan/test_all.py
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -34,12 +34,12 @@ def test_setup_data():
 
     assert isinstance(dataloader_train, DataLoader)
     assert isinstance(dataloader_eval, DataLoader)
-    train_batch = next(dataloader_train)
+    train_batch = next(iter(dataloader_train))
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
     assert isinstance(train_batch[1], Number)
     assert train_batch[0].ndim == 4
-    eval_batch = next(dataloader_eval)
+    eval_batch = next(iter(dataloader_eval))
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
     assert isinstance(eval_batch[1], Number)
@@ -65,5 +65,5 @@ def test_setup_trainer():
     trainer = setup_trainer(
         config, model, model, optimizer, optimizer, loss_fn, device
     )
-    trainer.run(batch)
+    trainer.run([batch, batch])
     assert isinstance(trainer.state.output, dict)
diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index b3acc52d..abf2b927 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -33,12 +33,12 @@ def test_setup_data():
 
     assert isinstance(dataloader_train, DataLoader)
     assert isinstance(dataloader_eval, DataLoader)
-    train_batch = next(dataloader_train)
+    train_batch = next(iter(dataloader_train))
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
     assert isinstance(train_batch[1], Number)
     assert train_batch[0].ndim == 4
-    eval_batch = next(dataloader_eval)
+    eval_batch = next(iter(dataloader_eval))
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
     assert isinstance(eval_batch[1], Number)
@@ -49,5 +49,5 @@ def test_setup_evaluator():
     model, _, device, _, batch = set_up()
     config = Namespace(use_amp=False)
     evaluator = setup_evaluator(config, model, device)
-    evaluator.run(batch)
+    evaluator.run([batch, batch])
     assert isinstance(evaluator.state.output, tuple)

From 608dd1d4bcb053682bd4c8c88a686a2a5b81be50 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 17:09:20 +0630
Subject: [PATCH 04/12] fix: test

---
 .../test_all.py                               |  7 ++++---
 .../template-vision-dcgan/test_all.py         |  6 ++++--
 .../template-vision-segmentation/data.py      |  2 +-
 .../template-vision-segmentation/test_all.py  | 20 +++++++++----------
 4 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/templates/template-vision-classification/test_all.py b/src/templates/template-vision-classification/test_all.py
index abf2b927..3313d3df 100644
--- a/src/templates/template-vision-classification/test_all.py
+++ b/src/templates/template-vision-classification/test_all.py
@@ -1,6 +1,5 @@
 import os
 from argparse import Namespace
-from numbers import Number
 from typing import Iterable
 
 import ignite.distributed as idist
@@ -36,13 +35,15 @@ def test_setup_data():
     train_batch = next(iter(dataloader_train))
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
-    assert isinstance(train_batch[1], Number)
+    assert isinstance(train_batch[1], Tensor)
     assert train_batch[0].ndim == 4
+    assert train_batch[1].ndim == 1
     eval_batch = next(iter(dataloader_eval))
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
-    assert isinstance(eval_batch[1], Number)
+    assert isinstance(eval_batch[1], Tensor)
     assert eval_batch[0].ndim == 4
+    assert eval_batch[1].ndim == 1
 
 
 def test_setup_evaluator():
diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
index 09a1cac8..94a9aa86 100644
--- a/src/templates/template-vision-dcgan/test_all.py
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -37,13 +37,15 @@ def test_setup_data():
     train_batch = next(iter(dataloader_train))
     assert isinstance(train_batch, Iterable)
     assert isinstance(train_batch[0], Tensor)
-    assert isinstance(train_batch[1], Number)
+    assert isinstance(train_batch[1], Tensor)
     assert train_batch[0].ndim == 4
+    assert train_batch[1].ndim == 1
     eval_batch = next(iter(dataloader_eval))
     assert isinstance(eval_batch, Iterable)
     assert isinstance(eval_batch[0], Tensor)
-    assert isinstance(eval_batch[1], Number)
+    assert isinstance(eval_batch[1], Tensor)
     assert eval_batch[0].ndim == 4
+    assert eval_batch[1].ndim == 4
 
 
 def test_models():
diff --git a/src/templates/template-vision-segmentation/data.py b/src/templates/template-vision-segmentation/data.py
index acd16add..79c403de 100644
--- a/src/templates/template-vision-segmentation/data.py
+++ b/src/templates/template-vision-segmentation/data.py
@@ -133,7 +133,7 @@ def setup_data(config: Namespace):
     dataloader_eval = idist.auto_dataloader(
         dataset_eval,
         shuffle=False,
-        batch_size=config.train_batch_size,
+        batch_size=config.eval_batch_size,
         num_workers=config.num_workers,
         drop_last=False,
     )
diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index abf2b927..688be54f 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -1,7 +1,5 @@
 import os
 from argparse import Namespace
-from numbers import Number
-from typing import Iterable
 
 import ignite.distributed as idist
 import pytest
@@ -34,15 +32,17 @@ def test_setup_data():
     assert isinstance(dataloader_train, DataLoader)
     assert isinstance(dataloader_eval, DataLoader)
     train_batch = next(iter(dataloader_train))
-    assert isinstance(train_batch, Iterable)
-    assert isinstance(train_batch[0], Tensor)
-    assert isinstance(train_batch[1], Number)
-    assert train_batch[0].ndim == 4
+    assert isinstance(train_batch, dict)
+    assert isinstance(train_batch["image"], Tensor)
+    assert isinstance(train_batch["mask"], Tensor)
+    assert train_batch["image"].ndim == 4
+    assert train_batch["mask"].ndim == 3
     eval_batch = next(iter(dataloader_eval))
-    assert isinstance(eval_batch, Iterable)
-    assert isinstance(eval_batch[0], Tensor)
-    assert isinstance(eval_batch[1], Number)
-    assert eval_batch[0].ndim == 4
+    assert isinstance(eval_batch, dict)
+    assert isinstance(eval_batch["image"], Tensor)
+    assert isinstance(eval_batch["mask"], Tensor)
+    assert eval_batch["image"].ndim == 4
+    assert eval_batch["mask"].ndim == 3
 
 
 def test_setup_evaluator():

From f76c4d5d62e8ef0c9f49044109cc153e2ead5a1e Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 17:15:12 +0630
Subject: [PATCH 05/12] fix: test

---
 src/templates/template-vision-dcgan/test_all.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
index 94a9aa86..6309d23d 100644
--- a/src/templates/template-vision-dcgan/test_all.py
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -7,7 +7,7 @@
 import pytest
 import torch
 from data import setup_data
-from models import Discriminator, Generator
+from model import Discriminator, Generator
 from torch import Tensor, nn, optim
 from torch.utils.data.dataloader import DataLoader
 from trainers import setup_trainer

From b82f540adb7a8145c95381d8539a795afb9ecaf9 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 17:22:31 +0630
Subject: [PATCH 06/12] fix: test

---
 src/templates/template-vision-dcgan/test_all.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
index 6309d23d..ec3be8ae 100644
--- a/src/templates/template-vision-dcgan/test_all.py
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -30,7 +30,7 @@ def test_setup_data():
     config = Namespace(
         data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
     )
-    dataloader_train, dataloader_eval = setup_data(config)
+    dataloader_train, dataloader_eval, _ = setup_data(config)
 
     assert isinstance(dataloader_train, DataLoader)
     assert isinstance(dataloader_eval, DataLoader)

From 1abe77be3ed419dbf65baec4547b1564f1fac094 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 17:26:50 +0630
Subject: [PATCH 07/12] fix: test

---
 src/templates/template-vision-dcgan/test_all.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
index ec3be8ae..8214f397 100644
--- a/src/templates/template-vision-dcgan/test_all.py
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -45,7 +45,7 @@ def test_setup_data():
     assert isinstance(eval_batch[0], Tensor)
     assert isinstance(eval_batch[1], Tensor)
     assert eval_batch[0].ndim == 4
-    assert eval_batch[1].ndim == 4
+    assert eval_batch[1].ndim == 1
 
 
 def test_models():

From aba8edd650d2a9513dfa62106f30feaa1c217a55 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 17:34:31 +0630
Subject: [PATCH 08/12] fix: test

---
 src/templates/template-vision-segmentation/test_all.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index 688be54f..1387e5f1 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -5,6 +5,7 @@
 import pytest
 import torch
 from data import setup_data
+from ignite.metrics import ConfusionMatrix, IoU
 from torch import Tensor, nn, optim
 from torch.utils.data.dataloader import DataLoader
 from trainers import setup_evaluator
@@ -48,6 +49,9 @@ def test_setup_data():
 def test_setup_evaluator():
     model, _, device, _, batch = set_up()
     config = Namespace(use_amp=False)
-    evaluator = setup_evaluator(config, model, device)
+    cm_metric = ConfusionMatrix(num_classes=21)
+    metrics = {"IoU": IoU(cm_metric)}
+
+    evaluator = setup_evaluator(config, model, metrics, device)
     evaluator.run([batch, batch])
     assert isinstance(evaluator.state.output, tuple)

From 797c2b2eeededdc794144e5825c00fbd4feb44ee Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 17:40:18 +0630
Subject: [PATCH 09/12] fix: test

---
 .../template-vision-segmentation/test_all.py        | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index 1387e5f1..83a68920 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -47,11 +47,18 @@ def test_setup_data():
 
 
 def test_setup_evaluator():
-    model, _, device, _, batch = set_up()
-    config = Namespace(use_amp=False)
+    model, _, device, _, _ = set_up()
+    config = Namespace(
+        use_amp=False,
+        data_path="~/data",
+        train_batch_size=1,
+        eval_batch_size=1,
+        num_workers=0,
+    )
+    dataloader_train, _ = setup_data(config)
     cm_metric = ConfusionMatrix(num_classes=21)
     metrics = {"IoU": IoU(cm_metric)}
 
     evaluator = setup_evaluator(config, model, metrics, device)
-    evaluator.run([batch, batch])
+    evaluator.run(dataloader_train)
     assert isinstance(evaluator.state.output, tuple)

From d056da7919820822237e3e361eddcc7e8ee324a1 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 18:01:54 +0630
Subject: [PATCH 10/12] fix: test

---
 src/templates/template-vision-segmentation/test_all.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index 83a68920..6746deeb 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -55,10 +55,10 @@ def test_setup_evaluator():
         eval_batch_size=1,
         num_workers=0,
     )
-    dataloader_train, _ = setup_data(config)
+    _, dataloader_eval = setup_data(config)
     cm_metric = ConfusionMatrix(num_classes=21)
     metrics = {"IoU": IoU(cm_metric)}
 
     evaluator = setup_evaluator(config, model, metrics, device)
-    evaluator.run(dataloader_train)
+    evaluator.run(dataloader_eval)
     assert isinstance(evaluator.state.output, tuple)

From 7e97aa3c70afcdb552fee88ceb787db0317e4bca Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 18:09:31 +0630
Subject: [PATCH 11/12] fix: test

---
 .../template-vision-segmentation/test_all.py  | 34 +------------------
 1 file changed, 1 insertion(+), 33 deletions(-)

diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
index 6746deeb..6859c879 100644
--- a/src/templates/template-vision-segmentation/test_all.py
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -1,24 +1,10 @@
 import os
 from argparse import Namespace
 
-import ignite.distributed as idist
 import pytest
-import torch
 from data import setup_data
-from ignite.metrics import ConfusionMatrix, IoU
-from torch import Tensor, nn, optim
+from torch import Tensor
 from torch.utils.data.dataloader import DataLoader
-from trainers import setup_evaluator
-
-
-def set_up():
-    model = nn.Linear(1, 1)
-    optimizer = optim.Adam(model.parameters())
-    device = idist.device()
-    loss_fn = nn.MSELoss()
-    batch = [torch.tensor([1.0]), torch.tensor([1.0])]
-
-    return model, optimizer, device, loss_fn, batch
 
 
 @pytest.mark.skipif(
@@ -44,21 +30,3 @@ def test_setup_data():
     assert isinstance(eval_batch["mask"], Tensor)
     assert eval_batch["image"].ndim == 4
     assert eval_batch["mask"].ndim == 3
-
-
-def test_setup_evaluator():
-    model, _, device, _, _ = set_up()
-    config = Namespace(
-        use_amp=False,
-        data_path="~/data",
-        train_batch_size=1,
-        eval_batch_size=1,
-        num_workers=0,
-    )
-    _, dataloader_eval = setup_data(config)
-    cm_metric = ConfusionMatrix(num_classes=21)
-    metrics = {"IoU": IoU(cm_metric)}
-
-    evaluator = setup_evaluator(config, model, metrics, device)
-    evaluator.run(dataloader_eval)
-    assert isinstance(evaluator.state.output, tuple)

From 8145f22ffbb54d8c4fbda8d75d9c57385950ed94 Mon Sep 17 00:00:00 2001
From: ydcjeff <32727188+ydcjeff@users.noreply.github.com>
Date: Tue, 25 May 2021 18:58:36 +0630
Subject: [PATCH 12/12] fix: delete test_all if uncheck

---
 src/store.js | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/store.js b/src/store.js
index 855afe14..4804dcd2 100644
--- a/src/store.js
+++ b/src/store.js
@@ -63,6 +63,7 @@ export function genCode() {
   if (currentFiles && Object.keys(currentFiles).length) {
     for (const file in currentFiles) {
       if (!store.config.include_test && file === 'test_all.py') {
+        delete store.code['test_all.py']
         continue
       }
       store.code[file] = ejs