From 1bfa3ffe678cf260aad247468c9ef41b3877322d Mon Sep 17 00:00:00 2001
From: Abhijit Deo <72816663+abhi-glitchhg@users.noreply.github.com>
Date: Wed, 13 Apr 2022 12:11:20 +0530
Subject: [PATCH] WIP

---
 configs/trial/vit_tiny.py |  6 ---
 tests/dir1/dir1_a.py      |  2 +
 tests/dir1/dir1_b.py      | 10 +++++
 tests/root_cfg.py         | 13 +++++++
 tests/test_config.py      | 78 ++++++++++++++++++++++++++++++++++++---
 5 files changed, 98 insertions(+), 11 deletions(-)
 delete mode 100644 configs/trial/vit_tiny.py
 create mode 100644 tests/dir1/dir1_a.py
 create mode 100644 tests/dir1/dir1_b.py
 create mode 100644 tests/root_cfg.py

diff --git a/configs/trial/vit_tiny.py b/configs/trial/vit_tiny.py
deleted file mode 100644
index 5f5e5274..00000000
--- a/configs/trial/vit_tiny.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from vformer.attention import VanillaSelfAttention
-from vformer.functional import PreNorm
-from vformer.config import LazyCall as L
-
-encoder = L(PreNorm(dim=96,
-                    fn=L(VanillaSelfAttention)(dim=96, num_heads=3, head_dim=96, p_dropout = 0.2)))
diff --git a/tests/dir1/dir1_a.py b/tests/dir1/dir1_a.py
new file mode 100644
index 00000000..cf2b5dff
--- /dev/null
+++ b/tests/dir1/dir1_a.py
@@ -0,0 +1,2 @@
+dir1a_str = "base_a_1"
+dir1a_dict = {"a": 1, "b": 2}
diff --git a/tests/dir1/dir1_b.py b/tests/dir1/dir1_b.py
new file mode 100644
index 00000000..0eb26a9a
--- /dev/null
+++ b/tests/dir1/dir1_b.py
@@ -0,0 +1,10 @@
+from vformer.config import LazyConfig
+
+# equivalent to relative import
+dir1a_str, dir1a_dict = LazyConfig.load_rel("dir1_a.py", ("dir1a_str", "dir1a_dict"))
+
+dir1b_str = dir1a_str + "_from_b"
+dir1b_dict = dir1a_dict
+
+# Every import is a reload: not modified by other config files
+assert dir1a_dict.a == 1
diff --git a/tests/root_cfg.py b/tests/root_cfg.py
new file mode 100644
index 00000000..139d5553
--- /dev/null
+++ b/tests/root_cfg.py
@@ -0,0 +1,13 @@
+from itertools import count
+
+from vformer.config import LazyCall as L
+
+from .dir1.dir1_a import dir1a_dict, dir1a_str
+
+dir1a_dict.a = "modified"
+
+# modification above won't affect future imports
+from .dir1.dir1_b import dir1b_dict, dir1b_str
+
+
+lazyobj = L(count)(x=dir1a_str, y=dir1b_str)
diff --git a/tests/test_config.py b/tests/test_config.py
index 56bdfd25..b50fcb1c 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,7 +1,14 @@
+import pytest
 import torch
 
 from vformer.config import LazyCall, instantiate, get_config
 from vformer.models import PVTSegmentation, SwinTransformer, VanillaViT, ViViTModel2
+import os
+import tempfile
+from itertools import count
+
+from vformer.config import LazyConfig, LazyCall as L
+from omegaconf import DictConfig
 
 
 def test_lazy():
@@ -48,11 +55,72 @@ def test_lazy():
     assert vivit(rand_vdo_tensor).shape == (32, 10)
 
 
-def test_lazyconfig():
-    file_addr = "../configs/trial/vit_tiny.py"
-    from vformer.config import LazyConfig
 
-    obj = LazyConfig()
+root_filename = os.path.join(os.path.dirname(__file__), "root_cfg.py")
+
+def test_load():
+    cfg = LazyConfig.load(root_filename)
+
+    assert cfg.dir1a_dict.a == "modified"
+
+    assert  cfg.dir1b_dict.a == 1
+    assert cfg.lazyobj.x == "base_a_1"
+
+    cfg.lazyobj.x = "new_x"
+    # reload
+    cfg = LazyConfig.load(root_filename)
+    assert  cfg.lazyobj.x == "base_a_1"
+
+def test_save_load():
+    cfg = LazyConfig.load(root_filename)
+    with tempfile.TemporaryDirectory(prefix="vformer") as d:
+        fname = os.path.join(d, "test_config.yaml")
+        LazyConfig.save(cfg, fname)
+        cfg2 = LazyConfig.load(fname)
+
+    assert cfg2.lazyobj._target_ == "itertools.count"
+    assert cfg.lazyobj._target_ == count
+    cfg2.lazyobj.pop("_target_")
+    cfg.lazyobj.pop("_target_")
+    # the rest are equal
+    assert cfg == cfg2
+
+def test_failed_save():
+    cfg = DictConfig({"x": lambda: 3}, flags={"allow_objects": True})
+    with tempfile.TemporaryDirectory(prefix="vformer") as d:
+        fname = os.path.join(d, "test_config.yaml")
+        LazyConfig.save(cfg, fname)
+        assert os.path.exists(fname) == True
+        assert  os.path.exists(fname + ".pkl") == True
+
+def test_overrides():
+    cfg = LazyConfig.load(root_filename)
+    LazyConfig.apply_overrides(cfg, ["lazyobj.x=123", 'dir1b_dict.a="123"'])
+    assert cfg.dir1b_dict.a == "123"
+    assert cfg.lazyobj.x == 123
 
-    print(LazyConfig.to_py(get_config(file_addr)))
+def test_invalid_overrides():
+    cfg = LazyConfig.load(root_filename)
+    with pytest.raises(KeyError):
+        LazyConfig.apply_overrides(cfg, ["lazyobj.x.xxx=123"])
 
+def test_to_py():
+    cfg = LazyConfig.load(root_filename)
+    cfg.lazyobj.x = {"a": 1, "b": 2, "c": L(count)(x={"r": "a", "s": 2.4, "t": [1, 2, 3, "z"]})}
+    cfg.list = ["a", 1, "b", 3.2]
+    py_str = LazyConfig.to_py(cfg)
+    expected = """cfg.dir1a_dict.a = "modified"
+cfg.dir1a_dict.b = 2
+cfg.dir1b_dict.a = 1
+cfg.dir1b_dict.b = 2
+cfg.lazyobj = itertools.count(
+x={
+    "a": 1,
+    "b": 2,
+    "c": itertools.count(x={"r": "a", "s": 2.4, "t": [1, 2, 3, "z"]}),
+},
+y="base_a_1_from_b",
+)
+cfg.list = ["a", 1, "b", 3.2]
+"""
+    assert py_str == expected