From 07a45109a71cd6f12d688961154bf8808863d65a Mon Sep 17 00:00:00 2001
From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Date: Tue, 6 Feb 2024 16:12:01 +0800
Subject: [PATCH] fix skipped tests (#3241)

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
---
 source/tests/pt/model/test_env_mat.py | 15 ++----
 source/tests/pt/model/test_mlp.py     | 75 +++++----------------------
 source/tests/pt/model/test_se_e2_a.py | 12 +----
 3 files changed, 16 insertions(+), 86 deletions(-)

diff --git a/source/tests/pt/model/test_env_mat.py b/source/tests/pt/model/test_env_mat.py
index b9f0ff1981..08890cc963 100644
--- a/source/tests/pt/model/test_env_mat.py
+++ b/source/tests/pt/model/test_env_mat.py
@@ -4,17 +4,9 @@
 import numpy as np
 import torch
 
-try:
-    from deepmd.dpmodel import (
-        EnvMat,
-    )
-
-    support_env_mat = True
-except ModuleNotFoundError:
-    support_env_mat = False
-except ImportError:
-    support_env_mat = False
-
+from deepmd.dpmodel.utils import (
+    EnvMat,
+)
 from deepmd.pt.model.descriptor.env_mat import (
     prod_env_mat_se_a,
 )
@@ -77,7 +69,6 @@ def setUp(self):
 
 
 # to be merged with the tf test case
-@unittest.skipIf(not support_env_mat, "EnvMat not supported")
 class TestEnvMat(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self):
         TestCaseSingleFrameWithNlist.setUp(self)
diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py
index 3a78b8294d..ca2bb6d2cf 100644
--- a/source/tests/pt/model/test_mlp.py
+++ b/source/tests/pt/model/test_mlp.py
@@ -5,6 +5,18 @@
 import numpy as np
 import torch
 
+from deepmd.dpmodel.utils import EmbeddingNet as DPEmbeddingNet
+from deepmd.dpmodel.utils import FittingNet as DPFittingNet
+from deepmd.dpmodel.utils import (
+    NativeLayer,
+    NativeNet,
+)
+from deepmd.pt.model.network.mlp import (
+    MLP,
+    EmbeddingNet,
+    FittingNet,
+    MLPLayer,
+)
 from deepmd.pt.utils import (
     env,
 )
@@ -12,65 +24,6 @@
     PRECISION_DICT,
 )
 
-try:
-    from deepmd.pt.model.network.mlp import (
-        MLP,
-        MLPLayer,
-    )
-
-    support_native_net = True
-except ModuleNotFoundError:
-    support_native_net = False
-
-try:
-    from deepmd.pt.model.network.mlp import (
-        EmbeddingNet,
-    )
-
-    support_embedding_net = True
-except ModuleNotFoundError:
-    support_embedding_net = False
-
-try:
-    from deepmd.pt.model.network.mlp import (
-        FittingNet,
-    )
-
-    support_fitting_net = True
-except ModuleNotFoundError:
-    support_fitting_net = False
-
-
-try:
-    from deepmd.dpmodel import (
-        NativeLayer,
-        NativeNet,
-    )
-
-    support_native_net = True
-except ModuleNotFoundError:
-    support_native_net = False
-except ImportError:
-    support_native_net = False
-
-try:
-    from deepmd.dpmodel import EmbeddingNet as DPEmbeddingNet
-
-    support_embedding_net = True
-except ModuleNotFoundError:
-    support_embedding_net = False
-except ImportError:
-    support_embedding_net = False
-
-try:
-    from deepmd.dpmodel import FittingNet as DPFittingNet
-
-    support_fitting_net = True
-except ModuleNotFoundError:
-    support_fitting_net = False
-except ImportError:
-    support_fitting_net = False
-
 
 def get_tols(prec):
     if prec in ["single", "float32"]:
@@ -84,7 +37,6 @@ def get_tols(prec):
     return rtol, atol
 
 
-@unittest.skipIf(not support_native_net, "NativeLayer not supported")
 class TestMLPLayer(unittest.TestCase):
     def setUp(self):
         self.test_cases = itertools.product(
@@ -141,7 +93,6 @@ def test_jit(self):
             model = torch.jit.script(ml1)
 
 
-@unittest.skipIf(not support_native_net, "NativeLayer not supported")
 class TestMLP(unittest.TestCase):
     def setUp(self):
         self.test_cases = itertools.product(
@@ -210,7 +161,6 @@ def test_jit(self):
             model = torch.jit.script(ml1)
 
 
-@unittest.skipIf(not support_embedding_net, "EmbeddingNet not supported")
 class TestEmbeddingNet(unittest.TestCase):
     def setUp(self):
         self.test_cases = itertools.product(
@@ -261,7 +211,6 @@ def test_jit(
             model = torch.jit.script(ml1)
 
 
-@unittest.skipIf(not support_fitting_net, "FittingNet not supported")
 class TestFittingNet(unittest.TestCase):
     def setUp(self):
         self.test_cases = itertools.product(
diff --git a/source/tests/pt/model/test_se_e2_a.py b/source/tests/pt/model/test_se_e2_a.py
index ec49725929..520fdfcdfa 100644
--- a/source/tests/pt/model/test_se_e2_a.py
+++ b/source/tests/pt/model/test_se_e2_a.py
@@ -5,16 +5,7 @@
 import numpy as np
 import torch
 
-try:
-    # from deepmd.dpmodel import PRECISION_DICT as DP_PRECISION_DICT
-    from deepmd.dpmodel import DescrptSeA as DPDescrptSeA
-
-    support_se_e2_a = True
-except ModuleNotFoundError:
-    support_se_e2_a = False
-except ImportError:
-    support_se_e2_a = False
-
+from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA
 from deepmd.pt.model.descriptor.se_a import (
     DescrptSeA,
 )
@@ -36,7 +27,6 @@
 
 
 # to be merged with the tf test case
-@unittest.skipIf(not support_se_e2_a, "EnvMat not supported")
 class TestDescrptSeA(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self):
         TestCaseSingleFrameWithNlist.setUp(self)