From c755731b1faa2e7b74e02386540b4b2aa546ab68 Mon Sep 17 00:00:00 2001
From: Simon Brugman <sbrugman@users.noreply.github.com>
Date: Thu, 5 Sep 2024 21:34:38 +0200
Subject: [PATCH] Fix distinct error codes test (#73)

Closes https://github.com/pytorch-labs/torchfix/issues/71
---
 tests/test_torchfix.py                       |  4 ++--
 torchfix/visitors/vision/singleton_import.py | 22 ++++----------------
 2 files changed, 6 insertions(+), 20 deletions(-)

diff --git a/tests/test_torchfix.py b/tests/test_torchfix.py
index 17ff183..3d9ddc2 100644
--- a/tests/test_torchfix.py
+++ b/tests/test_torchfix.py
@@ -65,8 +65,8 @@ def test_errorcodes_distinct():
         LOGGER.info("Checking error code for %s", visitor.__class__.__name__)
         errors = visitor.ERRORS
         for e in errors:
-            assert e not in seen
-            seen.add(e)
+            assert e.error_code not in seen
+            seen.add(e.error_code)
 
 
 def test_parse_error_code_str():
diff --git a/torchfix/visitors/vision/singleton_import.py b/torchfix/visitors/vision/singleton_import.py
index f64257b..f2b207b 100644
--- a/torchfix/visitors/vision/singleton_import.py
+++ b/torchfix/visitors/vision/singleton_import.py
@@ -9,22 +9,8 @@ class TorchVisionSingletonImportVisitor(TorchVisitor):
         TorchError(
             "TOR203",
             (
-                "Consider replacing 'import torchvision.datasets as datasets' "
-                "with 'from torchvision import datasets'."
-            ),
-        ),
-        TorchError(
-            "TOR203",
-            (
-                "Consider replacing 'import torchvision.models as models' "
-                "with 'from torchvision import models'."
-            ),
-        ),
-        TorchError(
-            "TOR203",
-            (
-                "Consider replacing 'import torchvision.transforms as transforms' "
-                "with 'from torchvision import transforms'."
+                "Consider replacing 'import torchvision.{module} as {module}' "
+                "with 'from torchvision import {module}'."
             ),
         ),
     ]
@@ -53,8 +39,8 @@ def visit_Import(self, node: cst.Import) -> None:
                         )
                     self.add_violation(
                         node,
-                        error_code=self.ERRORS[i].error_code,
-                        message=self.ERRORS[i].message(),
+                        error_code=self.ERRORS[0].error_code,
+                        message=self.ERRORS[0].message(module=import_attr),
                         replacement=replacement,
                     )
                 break