Repository: mlflow/mlflow

Issue:
Replace `unittest.TestCase.assertRaises` with `unittest.TestCase.assertRaisesRegex`
Some tests use `unittest.TestCase.assertRaises` to test an exception is raised for illegal operations, but they need to be replaces with `unittest.TestCase.assertRaisesRegex`.

### Why do we need this change?

Let's say we have a function that raises an exception:

```python
def throw_exception(...):
    if condition_1:
        raise TypeError("condition_1")
    if condition_2:
        raise TypeError("condition_2")
    ...
```

If we test this function using `assertRaises`:

```python
class MyTest(unittest.TestCase):
    def test_throw_exception(self):
        # Does `throw_exception` really raise the second TypeError?
        # It might throw the first TypeError, then the test will pass.
        with self.assertRaises(TypeError):
            throw_exception(...)  # should raise TypeError("condition_2")
```

If we test this function using `assertRaisesRegex`:

```python
class MyTest(unittest.TestCase):
    def test_throw_exception(self):
        # This test fails when `throw_exception` raises the first TypeError.
        with self. assertRaisesRegex(TypeError, "condition_b"):
            throw_exception(...)  # should raise TypeError("condition_2")
```

### Example

https://github.com/mlflow/mlflow/blob/fe6618823a2e6038149ee0da675503d2764552ca/tests/store/tracking/test_sqlalchemy_store.py#L107

The code above needs to be fixed to the following:

```python
 # "<string that matches the error message>" must be replaced
 with self.assertRaisesRegex(MlflowException, "<string that matches the error message>") as e:
```

### References

- https://docs.python.org/3/library/unittest.html#unittest.TestCase.assertRaises
- https://docs.python.org/3/library/unittest.html#unittest.TestCase.assertRaisesRegex

### Instructions

https://github.com/mlflow/mlflow/blob/101ad6e8eb383c769178df0df83d1d2a1cea6b4a/pylint_plugins/assert_raises_without_msg.py#L20-L33

Ping me with the file you want to work on :)

| File                                                  | Assignee    | PR    | Done |
| :---------------------------------------------------- | :---------- | :---- | :--- |
| `tests/entities/test_run_status.py`                   | @Sumanth077 |       |      |
| `tests/store/model_registry/test_sqlalchemy_store.py` | @ognis1205  | #5875 | ✅   |
| `tests/store/db/test_utils.py`                        | @erich-db   |       |      |
| `tests/store/tracking/__init__.py`                    | @Sumanth077 |       |      |
| `tests/store/tracking/test_file_store.py`             | @andy1122   |       |      |
| `tests/store/tracking/test_sqlalchemy_store.py`       | @ognis1205  | #5875 | ✅   |



Discussion:
cc @BenWilson2 
@erich-db this one is all yours!
@BenWilson2 On it!
Hi @erich-db, could you tell me which file you're going to fix?
@harupy I'll start with these three:
1. tests/entities/test_run_status.py
2. tests/store/db/test_utils.py
3. tests/store/tracking/__init__.py



@erich-db Sounds good! Looking forward to your PR!
@harupy I would like to work on this file.

 tests/store/tracking/test_file_store.py
@andy1122 Awesome, looking forward to your PR :)
@erich-db @andy1122 Looks like my understanding of the `msg` argument was wrong. Let me update the issue description.
Filed https://github.com/mlflow/mlflow/pull/5708 to fix the rule. @BenWilson2 Would you mind taking a look?
@harupy

I would like to contribute to the project so I am considering to participate in this issue to start with. It seems nobody is working on the following two files:

  tests/store/model_registry/test_sqlalchemy_store.py
  tests/store/tracking/test_sqlalchemy_store.py

Would you mind if I worked on these two files? Thank you.

@ognis1205 Thanks, looking forward to your PR!
@harupy I would like to contribute to this, do any other files remaining.
Thank You
@Sumanth077 Thanks! Could you fix `tests/entities/test_run_status.py` and `tests/store/tracking/__init__.py`?
Created the PR https://github.com/mlflow/mlflow/pull/5875
@harupy Would you mind reviewing it? Thank you.
@harupy Thankyou, will create a PR.



Resolving patch:
diff --git a/pylint_plugins/unittest_assert_raises.py b/pylint_plugins/unittest_assert_raises.py
--- a/pylint_plugins/unittest_assert_raises.py
+++ b/pylint_plugins/unittest_assert_raises.py
@@ -1,5 +1,3 @@
-import os
-
 import astroid
 from pylint.interfaces import IAstroidChecker
 from pylint.checkers import BaseChecker
@@ -9,31 +7,6 @@ def _is_unittest_assert_raises(node: astroid.Call):
     return isinstance(node.func, astroid.Attribute) and node.func.as_string() == "self.assertRaises"
 
 
-IGNORE_FILES = list(
-    map(
-        os.path.abspath,
-        [
-            # Instructions
-            # ============
-            # 1. Select a file in the list below and remove it.
-            # 2. Run pylint and confirm it fails.
-            # 3. Fix the lines printed out in the previous step.
-            # 4. Run pylint again and confirm it succeeds now.
-            # 5. Run pytest and confirm the changed lines don't fail.
-            # 6. Open a PR.
-            "tests/entities/test_run_status.py",
-            "tests/store/db/test_utils.py",
-            "tests/store/tracking/__init__.py",
-            "tests/store/tracking/test_file_store.py",
-        ],
-    )
-)
-
-
-def _should_ignore(path: str):
-    return path in IGNORE_FILES
-
-
 class UnittestAssertRaises(BaseChecker):
     __implements__ = IAstroidChecker
 
@@ -48,5 +21,5 @@ class UnittestAssertRaises(BaseChecker):
     priority = -1
 
     def visit_call(self, node: astroid.Call):
-        if not _should_ignore(node.root().file) and _is_unittest_assert_raises(node):
+        if _is_unittest_assert_raises(node):
             self.add_message(self.name, node=node)


Test patch:
diff --git a/tests/entities/test_run_status.py b/tests/entities/test_run_status.py
--- a/tests/entities/test_run_status.py
+++ b/tests/entities/test_run_status.py
@@ -34,13 +34,15 @@ def test_status_mappings(self):
         self.assertEqual("KILLED", RunStatus.to_string(RunStatus.KILLED))
         self.assertEqual(RunStatus.KILLED, RunStatus.from_string("KILLED"))
 
-        with self.assertRaises(Exception) as e:
+        with self.assertRaisesRegex(
+            Exception, r"Could not get string corresponding to run status -120"
+        ):
             RunStatus.to_string(-120)
-        self.assertIn("Could not get string corresponding to run status -120", str(e.exception))
 
-        with self.assertRaises(Exception) as e:
+        with self.assertRaisesRegex(
+            Exception, r"Could not get run status corresponding to string the IMPO"
+        ):
             RunStatus.from_string("the IMPOSSIBLE status string")
-        self.assertIn("Could not get run status corresponding to string the IMPO", str(e.exception))
 
     def test_is_terminated(self):
         self.assertTrue(RunStatus.is_terminated(RunStatus.FAILED))
diff --git a/tests/store/db/test_utils.py b/tests/store/db/test_utils.py
--- a/tests/store/db/test_utils.py
+++ b/tests/store/db/test_utils.py
@@ -61,20 +61,16 @@ def test_create_sqlalchemy_engine_with_retry_success_after_third_call(self):
                         self.assertEqual(engine, "Engine")
 
     def test_create_sqlalchemy_engine_with_retry_fail(self):
-        with mock.patch.dict(os.environ, {}):
-            with mock.patch("sqlalchemy.inspect") as mock_sqlalchemy_inspect:
-                with mock.patch(
-                    "mlflow.store.db.utils.create_sqlalchemy_engine"
-                ) as mock_create_sqlalchemy_engine:
-                    with mock.patch("time.sleep"):
-                        mock_sqlalchemy_inspect.side_effect = [Exception] * utils.MAX_RETRY_COUNT
-                        mock_create_sqlalchemy_engine.return_value = "Engine"
-                        self.assertRaises(
-                            Exception,
-                            utils.create_sqlalchemy_engine_with_retry,
-                            "mydb://host:port/",
-                        )
-                        assert (
-                            mock_create_sqlalchemy_engine.mock_calls
-                            == [mock.call("mydb://host:port/")] * utils.MAX_RETRY_COUNT
-                        )
+        with mock.patch.dict(os.environ, {}), mock.patch(
+            "sqlalchemy.inspect", side_effect=[Exception("failed")] * utils.MAX_RETRY_COUNT
+        ), mock.patch(
+            "mlflow.store.db.utils.create_sqlalchemy_engine", return_value="Engine"
+        ) as mock_create_sqlalchemy_engine, mock.patch(
+            "time.sleep"
+        ):
+            with self.assertRaisesRegex(Exception, r"failed"):
+                utils.create_sqlalchemy_engine_with_retry("mydb://host:port/")
+            assert (
+                mock_create_sqlalchemy_engine.mock_calls
+                == [mock.call("mydb://host:port/")] * utils.MAX_RETRY_COUNT
+            )
diff --git a/tests/store/tracking/__init__.py b/tests/store/tracking/__init__.py
--- a/tests/store/tracking/__init__.py
+++ b/tests/store/tracking/__init__.py
@@ -48,7 +48,9 @@ def test_record_logged_model(self):
                 RunTag(MLFLOW_LOGGED_MODELS, json.dumps([m.to_dict(), m2.to_dict(), m3.to_dict()]))
             ],
         )
-        with self.assertRaises(TypeError):
+        with self.assertRaisesRegex(
+            TypeError, "Argument 'mlflow_model' should be mlflow.models.Model, got '<class 'dict'>'"
+        ):
             store.record_logged_model(run_id, m.to_dict())
 
     @staticmethod
diff --git a/tests/store/tracking/test_file_store.py b/tests/store/tracking/test_file_store.py
--- a/tests/store/tracking/test_file_store.py
+++ b/tests/store/tracking/test_file_store.py
@@ -135,7 +135,7 @@ def test_valid_root(self):
         # Test removing root
         second_file_store = FileStore(self.test_root)
         shutil.rmtree(self.test_root)
-        with self.assertRaises(Exception):
+        with self.assertRaisesRegex(Exception, r"does not exist"):
             second_file_store._check_root_dir()
 
     def test_list_experiments(self):
@@ -176,8 +176,8 @@ def test_get_experiment(self):
         # test that fake experiments dont exist.
         # look for random experiment ids between 8000, 15000 since created ones are (100, 2000)
         for exp_id in set(random_int(8000, 15000) for x in range(20)):
-            with self.assertRaises(Exception):
-                fs.get_experiment(exp_id)
+            with self.assertRaisesRegex(Exception, f"Could not find experiment with ID {exp_id}"):
+                fs.get_experiment(str(exp_id))
 
     def test_get_experiment_int_experiment_id_backcompat(self):
         fs = FileStore(self.test_root)
@@ -214,9 +214,9 @@ def test_create_experiment(self):
         fs = FileStore(self.test_root)
 
         # Error cases
-        with self.assertRaises(Exception):
+        with self.assertRaisesRegex(Exception, "Invalid experiment name: 'None'"):
             fs.create_experiment(None)
-        with self.assertRaises(Exception):
+        with self.assertRaisesRegex(Exception, "Invalid experiment name: ''"):
             fs.create_experiment("")
 
         exp_id_ints = (int(exp_id) for exp_id in self.experiments)
@@ -296,7 +296,7 @@ def test_create_duplicate_experiments(self):
         fs = FileStore(self.test_root)
         for exp_id in self.experiments:
             name = self.exp_data[exp_id]["name"]
-            with self.assertRaises(Exception):
+            with self.assertRaisesRegex(Exception, f"Experiment '{name}' already exists"):
                 fs.create_experiment(name)
 
     def _extract_ids(self, experiments):
@@ -332,16 +332,17 @@ def test_rename_experiment(self):
         exp_id = self.experiments[random_int(0, len(self.experiments) - 1)]
 
         # Error cases
-        with self.assertRaises(Exception):
+        with self.assertRaisesRegex(Exception, "Invalid experiment name: 'None'"):
             fs.rename_experiment(exp_id, None)
-        with self.assertRaises(Exception):
-            # test that names of existing experiments are checked before renaming
-            other_exp_id = None
-            for exp in self.experiments:
-                if exp != exp_id:
-                    other_exp_id = exp
-                    break
-            fs.rename_experiment(exp_id, fs.get_experiment(other_exp_id).name)
+        # test that names of existing experiments are checked before renaming
+        other_exp_id = None
+        for exp in self.experiments:
+            if exp != exp_id:
+                other_exp_id = exp
+                break
+        name = fs.get_experiment(other_exp_id).name
+        with self.assertRaisesRegex(Exception, f"Experiment '{name}' already exists"):
+            fs.rename_experiment(exp_id, name)
 
         exp_name = self.exp_data[exp_id]["name"]
         new_name = exp_name + "!!!"
@@ -381,13 +382,13 @@ def test_hard_delete_run(self):
         exp_id = self.experiments[random_int(0, len(self.experiments) - 1)]
         run_id = self.exp_data[exp_id]["runs"][0]
         fs._hard_delete_run(run_id)
-        with self.assertRaises(MlflowException):
+        with self.assertRaisesRegex(MlflowException, f"Run '{run_id}' not found"):
             fs.get_run(run_id)
-        with self.assertRaises(MlflowException):
+        with self.assertRaisesRegex(MlflowException, f"Run '{run_id}' not found"):
             fs.get_all_tags(run_id)
-        with self.assertRaises(MlflowException):
+        with self.assertRaisesRegex(MlflowException, f"Run '{run_id}' not found"):
             fs.get_all_metrics(run_id)
-        with self.assertRaises(MlflowException):
+        with self.assertRaisesRegex(MlflowException, f"Run '{run_id}' not found"):
             fs.get_all_params(run_id)
 
     def test_get_deleted_runs(self):
@@ -669,9 +670,10 @@ def test_search_with_max_results(self):
         for n in [0, 1, 2, 4, 8, 10, 20, 50, 100, 500, 1000, 1200, 2000]:
             assert runs[: min(1200, n)] == self._search(fs, exp, max_results=n)
 
-        with self.assertRaises(MlflowException) as e:
+        with self.assertRaisesRegex(
+            MlflowException, "Invalid value for request parameter max_results. It "
+        ):
             self._search(fs, exp, None, max_results=int(1e10))
-        self.assertIn("Invalid value for request parameter max_results. It ", e.exception.message)
 
     def test_search_with_deterministic_max_results(self):
         fs = FileStore(self.test_root)
@@ -1013,18 +1015,16 @@ def _raise_exception_fn(*args, **kwargs):  # pylint: disable=unused-argument
             ]:
                 log_batch_kwargs = {"metrics": [], "params": [], "tags": []}
                 log_batch_kwargs.update(kwargs)
-                with self.assertRaises(MlflowException) as e:
+                with self.assertRaisesRegex(MlflowException, "Some internal error") as e:
                     fs.log_batch(run.info.run_id, **log_batch_kwargs)
-                self.assertIn(str(e.exception.message), "Some internal error")
                 assert e.exception.error_code == ErrorCode.Name(INTERNAL_ERROR)
 
     def test_log_batch_nonexistent_run(self):
         fs = FileStore(self.test_root)
         nonexistent_uuid = uuid.uuid4().hex
-        with self.assertRaises(MlflowException) as e:
+        with self.assertRaisesRegex(MlflowException, f"Run '{nonexistent_uuid}' not found") as e:
             fs.log_batch(nonexistent_uuid, [], [], [])
         assert e.exception.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
-        assert ("Run '%s' not found" % nonexistent_uuid) in e.exception.message
 
     def test_log_batch_params_idempotency(self):
         fs = FileStore(self.test_root)