Add unit, pipeline, case test

.github/workflows/docker.yml

@@ -7,9 +7,36 @@ on:
     - actions
 
 jobs:
+  unit_tests:
+    name: Unit Tests
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip'
+
+      - name: Install test dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install pytest
+
+      - name: Run unit tests
+        run: |
+          python -m pytest tests/unit -q
+
+      - name: Run pipeline smoke tests
+        run: |
+          python -m pytest tests/pipeline -q
+
   publish_package:
     name: Build and Publish
     runs-on: ubuntu-latest
+    needs: unit_tests
 
     steps:
       - uses: actions/checkout@v3

pyproject.toml

@@ -1,4 +1,10 @@
 [tool.pytest.ini_options]
-testpaths = ["tests"]
-addopts = "-ra"
-pythonpath = ["."]
+testpaths = ["tests/unit"]
+addopts = "-ra --strict-markers"
+pythonpath = [".", "tests/cases"]
+markers = [
+    "network: tests that require external network access",
+    "case: end-to-end usage cases that may train, predict, or download assets",
+    "integration: tests that exercise multiple subsystems together",
+    "slow: tests that are expected to take longer than unit tests",
+]

tests/README.md

@@ -0,0 +1,56 @@
+# Test Layout
+
+The test suite is split by purpose:
+
+- `unit/`: fast, deterministic tests that run locally without downloading data or model weights. This is the default pytest target.
+- `pipeline/`: offline smoke tests for the main `MolTrain.fit()` -> `MolPredict.predict()` workflow. These use fake data, tiny test dictionaries, random initialization, and one-epoch training to verify that the training and prediction pipeline is wired correctly without downloading pretrained weights.
+- `cases/`: end-to-end usage cases for training, prediction, representation extraction, and dataset/model downloads. These are not run by default.
+
+Run the default unit tests:
+
+```bash
+python -m pytest
+```
+
+Run the offline pipeline smoke tests:
+
+```bash
+python -m pytest tests/pipeline
+```
+
+Run both default CI test groups locally:
+
+```bash
+python -m pytest tests/unit tests/pipeline
+```
+
+Run the case tests without network access. Network-dependent tests will be collected but skipped:
+
+```bash
+python -m pytest tests/cases
+```
+
+Run network-dependent cases explicitly:
+
+```bash
+python -m pytest tests/cases --run-network
+```
+
+## Pipeline Smoke Tests
+
+Pipeline tests are intended to catch broken wiring across data loading, feature generation, model initialization, training, checkpoint saving, model reload, and prediction. They should stay small and deterministic:
+
+- Use local fake data instead of downloaded datasets.
+- Use `load_pretrained=False` so model weights are randomly initialized.
+- Provide a tiny local dictionary through `pretrained_dict_path` for UniMol v1 and conformer feature generation.
+- Use small batches, one epoch, CPU execution, and temporary output directories from pytest's `tmp_path`.
+- Assert shapes, finite predictions, and expected artifacts such as `model_0.pth`; do not assert model quality.
+
+The current pipeline smoke coverage should include both `unimolv1` and `unimolv2`, and both `regression` and `classification` tasks. Pretrained-weight behavior and dataset download behavior belong in `cases/` tests marked as `network` and/or `slow`, not in the default CI smoke path.
+
+Markers:
+
+- `network`: requires external network access.
+- `case`: user-facing end-to-end usage case.
+- `integration`: exercises multiple subsystems together.
+- `slow`: expected to take longer than unit tests.

tests/test_classification.py → tests/cases/test_classification.py

@@ -6,6 +6,8 @@
 
 from unimol_tools import MolTrain, MolPredict
 
+pytestmark = [pytest.mark.case, pytest.mark.integration, pytest.mark.slow]
+
 DATA_URL = 'https://weilab.math.msu.edu/DataLibrary/2D/Downloads/Ames_smi.zip'
 
 

tests/test_multiclass.py → tests/cases/test_multiclass.py

@@ -7,6 +7,8 @@
 
 from unimol_tools import MolTrain, MolPredict
 
+pytestmark = [pytest.mark.case, pytest.mark.integration, pytest.mark.slow]
+
 DATA_URL = 'https://weilab.math.msu.edu/DataLibrary/2D/Downloads/ESOL_smi.zip'
 
 

tests/test_multilabel_classification.py → tests/cases/test_multilabel_classification.py

@@ -8,6 +8,8 @@
 
 from unimol_tools import MolTrain, MolPredict
 
+pytestmark = [pytest.mark.case, pytest.mark.integration, pytest.mark.slow]
+
 CSV_URL = 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz'
 SDF_URL = 'https://tripod.nih.gov/tox21/challenge/download?id=tox21_10k_data_allsdf'
 

tests/test_multilabel_regression.py → tests/cases/test_multilabel_regression.py

@@ -7,6 +7,8 @@
 
 from unimol_tools import MolTrain, MolPredict
 
+pytestmark = [pytest.mark.case, pytest.mark.integration, pytest.mark.slow]
+
 DATA_URL = 'https://weilab.math.msu.edu/DataLibrary/2D/Downloads/FreeSolv_smi.zip'
 
 

tests/test_regression.py → tests/cases/test_regression.py

@@ -6,6 +6,8 @@
 
 from unimol_tools import MolTrain, MolPredict
 
+pytestmark = [pytest.mark.case, pytest.mark.integration, pytest.mark.slow]
+
 ESOL_TRAIN_URL = 'https://huggingface.co/datasets/HR-machine/ESol/resolve/main/train_data.csv?download=true'
 ESOL_TEST_URL = 'https://huggingface.co/datasets/HR-machine/ESol/resolve/main/test_data.csv?download=true'
 VQM24_URL = 'https://zenodo.org/records/15442257/files/DMC.npz?download=1'

tests/test_representation.py → tests/cases/test_representation.py

@@ -8,6 +8,8 @@
 
 from unimol_tools import UniMolRepr
 
+pytestmark = [pytest.mark.case, pytest.mark.integration, pytest.mark.slow]
+
 VQM24_URL = 'https://zenodo.org/records/15442257/files/DMC.npz?download=1'
 TOX21_CSV_URL = 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz'
 TOX21_SDF_URL = 'https://tripod.nih.gov/tox21/challenge/download?id=tox21_10k_data_allsdf'

tests/pipeline/test_smoke_pipeline.py

@@ -0,0 +1,136 @@
+import numpy as np
+import pytest
+
+from unimol_tools import MolPredict, MolTrain
+from unimol_tools.models import unimol as unimol_model_module
+from unimol_tools.models import unimolv2 as unimolv2_model_module
+
+
+pytestmark = [
+    pytest.mark.integration,
+    pytest.mark.filterwarnings(
+        "ignore:Precision is ill-defined and being set to 0.0 due to no predicted samples:sklearn.exceptions.UndefinedMetricWarning"
+    ),
+]
+
+
+def write_tiny_dictionary(tmp_path):
+    dict_path = tmp_path / "tiny.dict.txt"
+    dict_path.write_text("C 10\nH 10\nO 10\nN 10\n", encoding="utf-8")
+    return str(dict_path)
+
+
+def fake_pipeline_data(task):
+    target = (
+        [0.1, 0.2, 0.0, 0.3, 0.15, 0.25]
+        if task == "regression"
+        else [0, 1, 0, 1, 0, 1]
+    )
+    return {
+        "atoms": [
+            ["C", "O"],
+            ["C", "C"],
+            ["N", "C"],
+            ["O", "C"],
+            ["C", "N"],
+            ["C", "C", "O"],
+        ],
+        "coordinates": [
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
+            [[0.0, 0.0, 0.0], [1.2, 0.0, 0.0]],
+            [[0.0, 0.0, 0.0], [0.0, 1.1, 0.0]],
+            [[0.0, 0.0, 0.0], [0.0, 0.0, 1.1]],
+            [[0.0, 0.0, 0.0], [1.0, 1.0, 0.0]],
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.2, 0.0]],
+        ],
+        "target": target,
+    }
+
+
+def use_tiny_unimol_architectures(monkeypatch):
+    original_v1_architecture = unimol_model_module.molecule_architecture
+    original_v2_architecture = unimolv2_model_module.molecule_architecture
+
+    def tiny_v1_architecture():
+        args = original_v1_architecture()
+        args.encoder_layers = 1
+        args.encoder_embed_dim = 16
+        args.encoder_ffn_embed_dim = 32
+        args.encoder_attention_heads = 4
+        args.dropout = 0.0
+        args.emb_dropout = 0.0
+        args.attention_dropout = 0.0
+        args.max_seq_len = 32
+        return args
+
+    def tiny_v2_architecture(model_size="84m"):
+        args = original_v2_architecture(model_size)
+        args.num_encoder_layers = 1
+        args.encoder_embed_dim = 16
+        args.num_attention_heads = 4
+        args.encoder_attention_heads = 4
+        args.ffn_embedding_dim = 32
+        args.pair_embed_dim = 16
+        args.pair_hidden_dim = 8
+        args.dropout = 0.0
+        args.attention_dropout = 0.0
+        args.pair_dropout = 0.0
+        return args
+
+    monkeypatch.setattr(
+        unimol_model_module,
+        "molecule_architecture",
+        tiny_v1_architecture,
+    )
+    monkeypatch.setattr(
+        unimolv2_model_module,
+        "molecule_architecture",
+        tiny_v2_architecture,
+    )
+
+
+@pytest.mark.parametrize(
+    ("model_name", "task", "metrics"),
+    [
+        ("unimolv1", "regression", "mae"),
+        ("unimolv1", "classification", "acc"),
+        ("unimolv2", "regression", "mae"),
+        ("unimolv2", "classification", "acc"),
+    ],
+)
+def test_random_init_train_predict_pipeline(tmp_path, monkeypatch, model_name, task, metrics):
+    use_tiny_unimol_architectures(monkeypatch)
+    exp_dir = tmp_path / f"exp_{model_name}_{task}"
+    dict_path = write_tiny_dictionary(tmp_path)
+    train_data = fake_pipeline_data(task)
+    test_data = {
+        "atoms": train_data["atoms"][:2],
+        "coordinates": train_data["coordinates"][:2],
+    }
+
+    trainer = MolTrain(
+        task=task,
+        data_type="molecule",
+        epochs=1,
+        batch_size=2,
+        early_stopping=1,
+        kfold=2,
+        split="random",
+        metrics=metrics,
+        target_normalize="none",
+        use_cuda=False,
+        use_amp=False,
+        use_ddp=False,
+        model_name=model_name,
+        pretrained_dict_path=dict_path,
+        load_pretrained=False,
+        conf_cache_level=0,
+        save_path=str(exp_dir),
+    )
+
+    trainer.fit(train_data)
+    preds = MolPredict(load_model=str(exp_dir)).predict(test_data)
+
+    assert preds.shape == (len(test_data["atoms"]), 1)
+    assert np.isfinite(preds).all()
+    assert (exp_dir / "model_0.pth").exists()