format

Author: littleEast7

tests/local_cache/test_loadbalance.py

@@ -1,20 +1,25 @@
-import pytest
-from unittest.mock import patch
 import socket
+from unittest.mock import patch
+
+import pytest
 
-from alluxiofs.client.loadbalance import (
-    WorkerNetAddress,
-    WorkerListLoadBalancer,
-    DNSLoadBalancer,
-)
 from alluxiofs.client.config import AlluxioClientConfig
-from alluxiofs.client.const import ALLUXIO_WORKER_HTTP_SERVER_PORT_DEFAULT_VALUE
+from alluxiofs.client.const import (
+    ALLUXIO_WORKER_HTTP_SERVER_PORT_DEFAULT_VALUE,
+)
+from alluxiofs.client.loadbalance import DNSLoadBalancer
+from alluxiofs.client.loadbalance import WorkerListLoadBalancer
+from alluxiofs.client.loadbalance import WorkerNetAddress
+
 
 class TestWorkerNetAddress:
     def test_init_defaults(self):
         addr = WorkerNetAddress()
         assert addr.host == "localhost"
-        assert addr.http_server_port == ALLUXIO_WORKER_HTTP_SERVER_PORT_DEFAULT_VALUE
+        assert (
+            addr.http_server_port
+            == ALLUXIO_WORKER_HTTP_SERVER_PORT_DEFAULT_VALUE
+        )
 
     def test_from_host_and_port(self):
         addr = WorkerNetAddress.from_host_and_port("worker1", 12345)
@@ -24,7 +29,10 @@ def test_from_host_and_port(self):
     def test_from_host(self):
         addr = WorkerNetAddress.from_host("worker1")
         assert addr.host == "worker1"
-        assert addr.http_server_port == ALLUXIO_WORKER_HTTP_SERVER_PORT_DEFAULT_VALUE
+        assert (
+            addr.http_server_port
+            == ALLUXIO_WORKER_HTTP_SERVER_PORT_DEFAULT_VALUE
+        )
 
     def test_from_dict(self):
         d = {"host": "worker1", "http_server_port": 12345}
@@ -37,34 +45,45 @@ def test_to_dict(self):
         d = WorkerNetAddress.to_dict(addr)
         assert d == {"host": "worker1", "http_server_port": 12345}
 
+
 class TestWorkerListLoadBalancer:
     def test_init_hosts_only(self):
-        config = AlluxioClientConfig(worker_hosts="worker1,worker2", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            worker_hosts="worker1,worker2", worker_http_port=29999
+        )
         lb = WorkerListLoadBalancer(config)
         assert lb.workers == ["worker1", "worker2"]
         assert lb.ports == [29999, 29999]
 
     def test_init_hosts_with_ports(self):
-        config = AlluxioClientConfig(worker_hosts="worker1:10001,worker2:10002")
+        config = AlluxioClientConfig(
+            worker_hosts="worker1:10001,worker2:10002"
+        )
         lb = WorkerListLoadBalancer(config)
         assert lb.workers == ["worker1", "worker2"]
         assert lb.ports == [10001, 10002]
 
     def test_init_mixed(self):
-        config = AlluxioClientConfig(worker_hosts="worker1,worker2:10002", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            worker_hosts="worker1,worker2:10002", worker_http_port=29999
+        )
         lb = WorkerListLoadBalancer(config)
         assert lb.workers == ["worker1", "worker2"]
         assert lb.ports == [29999, 10002]
 
     def test_get_worker_random(self):
-        config = AlluxioClientConfig(worker_hosts="worker1", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            worker_hosts="worker1", worker_http_port=29999
+        )
         lb = WorkerListLoadBalancer(config)
         worker = lb.get_worker()
         assert worker.host == "worker1"
         assert worker.http_server_port == 29999
 
     def test_get_worker_hash(self):
-        config = AlluxioClientConfig(worker_hosts="worker1,worker2", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            worker_hosts="worker1,worker2", worker_http_port=29999
+        )
         lb = WorkerListLoadBalancer(config)
         path = "test_path"
         worker1 = lb.get_worker(path)
@@ -73,25 +92,30 @@ def test_get_worker_hash(self):
         assert worker1.host in ["worker1", "worker2"]
 
     def test_get_multiple_worker(self):
-        config = AlluxioClientConfig(worker_hosts="worker1,worker2,worker3", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            worker_hosts="worker1,worker2,worker3", worker_http_port=29999
+        )
         lb = WorkerListLoadBalancer(config)
         workers = lb.get_multiple_worker("test_path", 2)
         assert len(workers) == 2
         assert workers[0] != workers[1]
         assert workers[0].host in ["worker1", "worker2", "worker3"]
         assert workers[1].host in ["worker1", "worker2", "worker3"]
 
+
 class TestDNSLoadBalancer:
     @patch("socket.getaddrinfo")
     def test_resolve_address_success(self, mock_getaddrinfo):
         mock_getaddrinfo.return_value = [
-            (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('192.168.1.1', 80))
+            (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("192.168.1.1", 80))
         ]
         config = AlluxioClientConfig(load_balance_domain="example.com")
         lb = DNSLoadBalancer(config)
         ip = lb.resolve_address()
         assert ip == "192.168.1.1"
-        mock_getaddrinfo.assert_called_with("example.com", None, family=socket.AF_INET)
+        mock_getaddrinfo.assert_called_with(
+            "example.com", None, family=socket.AF_INET
+        )
 
     @patch("socket.getaddrinfo")
     def test_resolve_address_no_address(self, mock_getaddrinfo):
@@ -112,7 +136,9 @@ def test_resolve_address_gaierror(self, mock_getaddrinfo):
     @patch("alluxiofs.client.loadbalance.DNSLoadBalancer.resolve_address")
     def test_get_worker(self, mock_resolve):
         mock_resolve.return_value = "192.168.1.1"
-        config = AlluxioClientConfig(load_balance_domain="example.com", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            load_balance_domain="example.com", worker_http_port=29999
+        )
         lb = DNSLoadBalancer(config)
         worker = lb.get_worker("path")
         assert worker.host == "192.168.1.1"
@@ -121,7 +147,9 @@ def test_get_worker(self, mock_resolve):
     @patch("alluxiofs.client.loadbalance.DNSLoadBalancer.resolve_address")
     def test_get_worker_no_args(self, mock_resolve):
         mock_resolve.return_value = "192.168.1.1"
-        config = AlluxioClientConfig(load_balance_domain="example.com", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            load_balance_domain="example.com", worker_http_port=29999
+        )
         lb = DNSLoadBalancer(config)
         worker = lb.get_worker()
         assert worker.host == "192.168.1.1"
@@ -130,7 +158,9 @@ def test_get_worker_no_args(self, mock_resolve):
     @patch("alluxiofs.client.loadbalance.DNSLoadBalancer.resolve_address")
     def test_get_multiple_worker(self, mock_resolve):
         mock_resolve.side_effect = ["192.168.1.1", "192.168.1.2"]
-        config = AlluxioClientConfig(load_balance_domain="example.com", worker_http_port=29999)
+        config = AlluxioClientConfig(
+            load_balance_domain="example.com", worker_http_port=29999
+        )
         lb = DNSLoadBalancer(config)
         workers = lb.get_multiple_worker("path", 2)
         assert len(workers) == 2

tests/local_cache/test_local_cache.py

@@ -70,6 +70,7 @@ def test_atomic_write_success(self, mock_fetch):
         file_path = "/test/file"
         part_index = 0
         path_hash = self.manager.compute_hash(file_path)
+        assert path_hash is not None
         local_path = self.manager._get_local_path(file_path, part_index)
 
         # Mock fetch to write some data
@@ -225,7 +226,7 @@ def test_set_file_loading(self):
         os.remove(local_path + "_loading")
 
         # Fail (already cached)
-        with open(local_path, "w") as f:
+        with open(local_path, "w"):
             pass
         self.assertFalse(self.manager.set_file_loading(file_path, part_index))
 
@@ -292,6 +293,7 @@ def test_read_file_range_miss_triggers_download(self, mock_download):
         # We need to mock manager.read_from_cache because it's called inside read_file_range
 
         original_read = self.manager.read_from_cache
+        assert original_read is not None
 
         call_count = 0
 

tests/local_cache/test_prefetch_policy.py

@@ -1,9 +1,9 @@
 from unittest.mock import MagicMock
-from alluxiofs.client.prefetch_policy import (
-    NoPrefetchPolicy,
-    FixedWindowPrefetchPolicy,
-    AdaptiveWindowPrefetchPolicy,
-)
+
+from alluxiofs.client.prefetch_policy import AdaptiveWindowPrefetchPolicy
+from alluxiofs.client.prefetch_policy import FixedWindowPrefetchPolicy
+from alluxiofs.client.prefetch_policy import NoPrefetchPolicy
+
 
 class TestNoPrefetchPolicy:
     def test_get_blocks(self):
@@ -23,6 +23,7 @@ def test_get_blocks_full_file(self):
         # file_size=5000 -> blocks 0, 1, 2, 3, 4. (4*1024=4096, 5000 is in block 4)
         assert policy.get_blocks(0, -1, 5000) == (0, 4)
 
+
 class TestFixedWindowPrefetchPolicy:
     def test_init_defaults(self):
         policy = FixedWindowPrefetchPolicy(block_size=1024)
@@ -67,6 +68,7 @@ def test_get_blocks_boundary(self):
         # Should return (0, 4)
         assert policy.get_blocks(0, 1024, 5000) == (0, 4)
 
+
 class TestAdaptiveWindowPrefetchPolicy:
     def test_init(self):
         config = MagicMock()
@@ -122,7 +124,7 @@ def test_random_access_resets_window(self):
 
         policy = AdaptiveWindowPrefetchPolicy(block_size=1024, config=config)
         policy.prefetch_ahead = 5
-        policy.last_offset = 10240 # block 10
+        policy.last_offset = 10240  # block 10
 
         # Jump back to block 0
         # index=0, last_index=10. index < last_index.
@@ -139,11 +141,10 @@ def test_jump_forward_resets_window(self):
 
         policy = AdaptiveWindowPrefetchPolicy(block_size=1024, config=config)
         policy.prefetch_ahead = 2
-        policy.last_offset = 0 # block 0
+        policy.last_offset = 0  # block 0
 
         # Jump far forward to block 10
         # index=10, last_index=0. index-last_index-1 = 9 > prefetch_ahead(2).
         # else branch: prefetch_ahead = min_prefetch (0).
         policy.get_blocks(10240, 1024, 100000)
         assert policy.prefetch_ahead == 0
-

tests/local_cache/test_transfer.py

@@ -1,9 +1,13 @@
-import pytest
 import json
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+import pytest
 
-from alluxiofs.client.transfer import UfsInfo, UFSUpdater
 from alluxiofs.client.const import ALLUXIO_UFS_INFO_REFRESH_INTERVAL_MINUTES
+from alluxiofs.client.transfer import UfsInfo
+from alluxiofs.client.transfer import UFSUpdater
+
 
 class TestUfsInfo:
     def test_init(self):
@@ -12,6 +16,7 @@ def test_init(self):
         assert info.ufs_full_path == "ufs_path"
         assert info.options == {"opt": "val"}
 
+
 class TestUFSUpdater:
     @pytest.fixture
     def mock_alluxio(self):
@@ -31,17 +36,17 @@ def test_init_with_alluxio(self, mock_alluxio):
 
     def test_init_without_alluxio(self):
         updater = UFSUpdater(None)
-        assert updater.interval_seconds == ALLUXIO_UFS_INFO_REFRESH_INTERVAL_MINUTES * 60
+        assert (
+            updater.interval_seconds
+            == ALLUXIO_UFS_INFO_REFRESH_INTERVAL_MINUTES * 60
+        )
         assert updater.alluxio is None
 
     def test_parse_ufs_info_valid(self, mock_alluxio):
         updater = UFSUpdater(mock_alluxio)
-        json_str = json.dumps({
-            "s3://bucket/path": {
-                "alluxio_path": "/mnt/s3",
-                "key": "value"
-            }
-        })
+        json_str = json.dumps(
+            {"s3://bucket/path": {"alluxio_path": "/mnt/s3", "key": "value"}}
+        )
         result = updater.parse_ufs_info(json_str)
         assert len(result) == 1
         assert result[0].ufs_full_path == "s3://bucket/path"
@@ -73,12 +78,19 @@ def test_parse_ufs_info_value_not_dict(self, mock_alluxio):
     @patch("alluxiofs.client.transfer.register_unregistered_ufs_to_fsspec")
     @patch("alluxiofs.client.transfer.filesystem")
     @patch("alluxiofs.client.transfer.fsspec.get_filesystem_class")
-    def test_register_ufs_fallback(self, mock_get_fs_class, mock_filesystem, mock_register, mock_get_protocol, mock_alluxio):
+    def test_register_ufs_fallback(
+        self,
+        mock_get_fs_class,
+        mock_filesystem,
+        mock_register,
+        mock_get_protocol,
+        mock_alluxio,
+    ):
         updater = UFSUpdater(mock_alluxio)
         ufs_info = UfsInfo("/mnt/s3", "s3://bucket", {"key": "val"})
 
         mock_get_protocol.return_value = "s3"
-        mock_get_fs_class.return_value = True # Simulate supported protocol
+        mock_get_fs_class.return_value = True  # Simulate supported protocol
         mock_fs_instance = MagicMock()
         mock_filesystem.return_value = mock_fs_instance
 
@@ -92,11 +104,13 @@ def test_register_ufs_fallback(self, mock_get_fs_class, mock_filesystem, mock_re
 
     @patch("alluxiofs.client.transfer.register_unregistered_ufs_to_fsspec")
     @patch("alluxiofs.client.transfer.fsspec.get_filesystem_class")
-    def test_register_ufs_fallback_unsupported(self, mock_get_fs_class, mock_register, mock_alluxio):
+    def test_register_ufs_fallback_unsupported(
+        self, mock_get_fs_class, mock_register, mock_alluxio
+    ):
         updater = UFSUpdater(mock_alluxio)
         ufs_info = UfsInfo("/mnt/unknown", "unknown://bucket", {})
 
-        mock_get_fs_class.return_value = None # Simulate unsupported protocol
+        mock_get_fs_class.return_value = None  # Simulate unsupported protocol
 
         with pytest.raises(ValueError, match="Unsupported protocol"):
             updater.register_ufs_fallback([ufs_info])
@@ -113,12 +127,20 @@ def test_get_alluxio_path_from_ufs_full_path(self, mock_alluxio):
         updater._cached_ufs = {"s3://bucket": "fs_instance"}
         updater._path_map = {"s3://bucket": "/mnt/s3"}
 
-        assert updater.get_alluxio_path_from_ufs_full_path("s3://bucket/file") == "/mnt/s3/file"
-        assert updater.get_alluxio_path_from_ufs_full_path("hdfs://namenode/file") is None
+        assert (
+            updater.get_alluxio_path_from_ufs_full_path("s3://bucket/file")
+            == "/mnt/s3/file"
+        )
+        assert (
+            updater.get_alluxio_path_from_ufs_full_path("hdfs://namenode/file")
+            is None
+        )
 
     @patch("alluxiofs.client.transfer.UFSUpdater.parse_ufs_info")
     @patch("alluxiofs.client.transfer.UFSUpdater.register_ufs_fallback")
-    def test_execute_update_success(self, mock_register, mock_parse, mock_alluxio):
+    def test_execute_update_success(
+        self, mock_register, mock_parse, mock_alluxio
+    ):
         updater = UFSUpdater(mock_alluxio)
         mock_alluxio.get_ufs_info_from_worker.return_value = '{"json": "data"}'
         mock_parse.return_value = ["parsed_info"]
@@ -139,7 +161,9 @@ def test_execute_update_fail_none(self, mock_alluxio):
 
     def test_execute_update_exception(self, mock_alluxio):
         updater = UFSUpdater(mock_alluxio)
-        mock_alluxio.get_ufs_info_from_worker.side_effect = Exception("Network error")
+        mock_alluxio.get_ufs_info_from_worker.side_effect = Exception(
+            "Network error"
+        )
 
         updater._execute_update()
         # Should catch exception and log error
@@ -186,6 +210,9 @@ def test_must_get_methods_wait(self, mock_alluxio):
         with pytest.raises(ValueError):
             updater.must_get_ufs_from_path("path")
 
-        updater.get_alluxio_path_from_ufs_full_path = MagicMock(return_value="/path")
-        assert updater.must_get_alluxio_path_from_ufs_full_path("path") == "/path"
-
+        updater.get_alluxio_path_from_ufs_full_path = MagicMock(
+            return_value="/path"
+        )
+        assert (
+            updater.must_get_alluxio_path_from_ufs_full_path("path") == "/path"
+        )