Coverage for src / piccione / upload / cache_manager.py: 100%

1from typing import Set 

2 

3import redis 

4from redis.exceptions import ConnectionError as RedisConnectionError 

5 

6 

7class CacheManager: 

8 REDIS_KEY = "processed_files" 

9 

10 def __init__( 

11 self, 

12 redis_host: str = "localhost", 

13 redis_port: int = 6379, 

14 redis_db: int = 4, 

15 ): 

16 self._redis = None 

17 self.redis_host = redis_host 

18 self.redis_port = redis_port 

19 self.redis_db = redis_db 

20 self.processed_files: Set[str] = set() 

21 

22 self._init_cache() 

23 

24 def _init_redis(self) -> None: 

25 """Initialize Redis connection.""" 

26 try: 

27 self._redis = redis.Redis( 

28 host=self.redis_host, 

29 port=self.redis_port, 

30 db=self.redis_db, 

31 decode_responses=True, 

32 ) 

33 self._redis.ping() 

34 except RedisConnectionError: 

35 raise RuntimeError("Redis is not available. Cache requires Redis.") 

36 

37 def _init_cache(self) -> None: 

38 """Initialize cache from Redis.""" 

39 self._init_redis() 

40 existing_redis_files = self._redis.smembers(self.REDIS_KEY) 

41 self.processed_files.update(existing_redis_files) 

42 

43 def add(self, filename: str) -> None: 

44 """Add a file to the cache.""" 

45 self.processed_files.add(filename) 

46 self._redis.sadd(self.REDIS_KEY, filename) 

47 

48 def __contains__(self, filename: str) -> bool: 

49 """Check if a file is in the cache.""" 

50 return filename in self.processed_files 

51 

52 def get_all(self) -> Set[str]: 

53 """Return all files in the cache.""" 

54 self.processed_files.update(self._redis.smembers(self.REDIS_KEY)) 

55 return self.processed_files