Allow to set-credentials to all hosts (old style) (#76)

vrslev · Aug 29, 2024 · 3400138 · 3400138
1 parent 81a85c1
commit 3400138
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 13 deletions.
diff --git a/stompman/config.py b/stompman/config.py
@@ -36,7 +36,7 @@ def unescaped_passcode(self) -> str:
         return unquote(self.passcode)
 
     @classmethod
-    def from_pydantic_multihost_hosts(cls, hosts: list[MultiHostHostLike]) -> list[Self]:
+    def from_pydantic_multihost_hosts(cls, hosts: list[MultiHostHostLike]) -> list[Self]:  # noqa: C901
         """Create connection parameters from `pydantic_code.MultiHostUrl.hosts()`.
 
         .. code-block:: python
@@ -52,7 +52,8 @@ def from_pydantic_multihost_hosts(cls, hosts: list[MultiHostHostLike]) -> list[S
 
             async with stompman.Client(
                 servers=stompman.ConnectionParameters.from_pydantic_multihost_hosts(
-                    ArtemisDsn("tcp://lev:pass@host1:61616,host2:61617,host3:61618").hosts()
+                    ArtemisDsn("tcp://user:pass@host1:61616,host2:61617,host3:61618").hosts()
+                    # or: ArtemisDsn("tcp://user1:pass1@host1:61616,user2:pass2@host2:61617,user3:pass@host3:61618").hosts()
                 ),
             ):
                 ...
@@ -81,12 +82,18 @@ def from_pydantic_multihost_hosts(cls, hosts: list[MultiHostHostLike]) -> list[S
             else:
                 all_credentials.append((username, password))
 
-        if not all_credentials:
-            msg = "username and password must be set"
-            raise ValueError(msg)
-        if len(all_credentials) != 1:
-            msg = "only one username-password pair must be set"
-            raise ValueError(msg)
-
-        login, passcode = all_credentials[0]
-        return [cls(host=host, port=port, login=login, passcode=passcode) for (host, port) in all_hosts]
+        match len(all_credentials):
+            case value if value == len(all_hosts):
+                return [
+                    cls(host=host, port=port, login=username, passcode=password)
+                    for ((host, port), (username, password)) in zip(all_hosts, all_credentials, strict=True)
+                ]
+            case 1:
+                username, password = all_credentials[0]
+                return [cls(host=host, port=port, login=username, passcode=password) for (host, port) in all_hosts]
+            case 0:
+                msg = "username and password must be set"
+                raise ValueError(msg)
+            case _:
+                msg = "all username-password pairs or only one pair must be set"
+                raise ValueError(msg)
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -22,7 +22,7 @@ class MultiHostHostLikeFactory(TypedDictFactory[MultiHostHostLike]): ...
 
 
 class TestConnectionParametersFromPydanticMultiHostHosts:
-    def test_ok(self, faker: faker.Faker) -> None:
+    def test_ok_with_one_credentials(self, faker: faker.Faker) -> None:
         hosts: list[MultiHostHostLike] = [
             {"host": "host1", "port": 1, "username": None, "password": None},
             {"host": "host2", "port": 2, "username": None, "password": None},
@@ -42,6 +42,23 @@ def test_ok(self, faker: faker.Faker) -> None:
             stompman.ConnectionParameters("host4", 4, "lev", "pass"),
         ]
 
+    def test_ok_with_all_credentials(self, faker: faker.Faker) -> None:
+        hosts: list[MultiHostHostLike] = [
+            {"host": "host1", "port": 1, "username": "user1", "password": "pass1"},
+            {"host": "host2", "port": 2, "username": "user2", "password": "pass2"},
+            {"host": "host3", "port": 3, "username": "user3", "password": "pass3"},
+            {"host": "host4", "port": 4, "username": "user4", "password": "pass4"},
+        ]
+
+        result = stompman.ConnectionParameters.from_pydantic_multihost_hosts(hosts)
+
+        assert result == [
+            stompman.ConnectionParameters("host1", 1, "user1", "pass1"),
+            stompman.ConnectionParameters("host2", 2, "user2", "pass2"),
+            stompman.ConnectionParameters("host3", 3, "user3", "pass3"),
+            stompman.ConnectionParameters("host4", 4, "user4", "pass4"),
+        ]
+
     def test_no_host_or_port_or_both(self, faker: faker.Faker) -> None:
         cases: list[MultiHostHostLike] = [
             {"host": None, "port": faker.pyint(), "username": faker.pystr(), "password": faker.pystr()},
@@ -88,5 +105,5 @@ def test_multiple_credentials(self, faker: faker.Faker) -> None:
             {"host": faker.pystr(), "port": faker.pyint(), "username": faker.pystr(), "password": faker.pystr()},
         ]
 
-        with pytest.raises(ValueError, match="only one username-password pair must be set"):
+        with pytest.raises(ValueError, match="all username-password pairs or only one pair must be set"):
             stompman.ConnectionParameters.from_pydantic_multihost_hosts(hosts)