Coverage for src / dataknobs_bots / registry / memory.py: 30%

1"""In-memory implementation of RegistryBackend.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from datetime import datetime, timezone 

7from typing import Any 

8 

9from .models import Registration 

10 

11 

12class InMemoryBackend: 

13 """In-memory implementation of RegistryBackend. 

14 

15 Simple dict-based storage suitable for: 

16 - Testing without database dependencies 

17 - Single-instance deployments 

18 - Development environments 

19 

20 Not suitable for: 

21 - Multi-instance deployments (no persistence) 

22 - Production with persistence requirements 

23 

24 Thread-safety is provided via asyncio.Lock. 

25 

26 Example: 

27 ```python 

28 backend = InMemoryBackend() 

29 await backend.initialize() 

30 

31 reg = await backend.register("my-bot", {"llm": {...}}) 

32 print(f"Created at: {reg.created_at}") 

33 

34 config = await backend.get_config("my-bot") 

35 print(f"Config: {config}") 

36 

37 # List all bots 

38 for reg in await backend.list_active(): 

39 print(f"Bot: {reg.bot_id}") 

40 

41 # Cleanup 

42 await backend.close() 

43 ``` 

44 """ 

45 

46 def __init__(self) -> None: 

47 """Initialize the in-memory backend.""" 

48 self._registrations: dict[str, Registration] = {} 

49 self._lock = asyncio.Lock() 

50 self._initialized = False 

51 

52 async def initialize(self) -> None: 

53 """Initialize the backend (no-op for in-memory).""" 

54 self._initialized = True 

55 

56 async def close(self) -> None: 

57 """Close the backend (clears all data).""" 

58 async with self._lock: 

59 self._registrations.clear() 

60 self._initialized = False 

61 

62 async def register( 

63 self, 

64 bot_id: str, 

65 config: dict[str, Any], 

66 status: str = "active", 

67 ) -> Registration: 

68 """Register or update a bot. 

69 

70 Args: 

71 bot_id: Unique bot identifier 

72 config: Bot configuration dictionary 

73 status: Registration status (default: active) 

74 

75 Returns: 

76 Registration object with metadata 

77 """ 

78 async with self._lock: 

79 now = datetime.now(timezone.utc) 

80 

81 if bot_id in self._registrations: 

82 # Update existing - preserve created_at 

83 old = self._registrations[bot_id] 

84 reg = Registration( 

85 bot_id=bot_id, 

86 config=config, 

87 status=status, 

88 created_at=old.created_at, 

89 updated_at=now, 

90 last_accessed_at=now, 

91 ) 

92 else: 

93 # Create new 

94 reg = Registration( 

95 bot_id=bot_id, 

96 config=config, 

97 status=status, 

98 created_at=now, 

99 updated_at=now, 

100 last_accessed_at=now, 

101 ) 

102 

103 self._registrations[bot_id] = reg 

104 return reg 

105 

106 async def get(self, bot_id: str) -> Registration | None: 

107 """Get registration and update access time. 

108 

109 Args: 

110 bot_id: Bot identifier 

111 

112 Returns: 

113 Registration if found, None otherwise 

114 """ 

115 async with self._lock: 

116 reg = self._registrations.get(bot_id) 

117 if reg: 

118 # Update access time 

119 self._registrations[bot_id] = Registration( 

120 bot_id=reg.bot_id, 

121 config=reg.config, 

122 status=reg.status, 

123 created_at=reg.created_at, 

124 updated_at=reg.updated_at, 

125 last_accessed_at=datetime.now(timezone.utc), 

126 ) 

127 return self._registrations[bot_id] 

128 return None 

129 

130 async def get_config(self, bot_id: str) -> dict[str, Any] | None: 

131 """Get just the config. 

132 

133 Args: 

134 bot_id: Bot identifier 

135 

136 Returns: 

137 Config dict if found, None otherwise 

138 """ 

139 reg = await self.get(bot_id) 

140 return reg.config if reg else None 

141 

142 async def exists(self, bot_id: str) -> bool: 

143 """Check if active registration exists. 

144 

145 Args: 

146 bot_id: Bot identifier 

147 

148 Returns: 

149 True if registration exists and is active 

150 """ 

151 async with self._lock: 

152 reg = self._registrations.get(bot_id) 

153 return reg is not None and reg.status == "active" 

154 

155 async def unregister(self, bot_id: str) -> bool: 

156 """Hard delete registration. 

157 

158 Args: 

159 bot_id: Bot identifier 

160 

161 Returns: 

162 True if deleted, False if not found 

163 """ 

164 async with self._lock: 

165 if bot_id in self._registrations: 

166 del self._registrations[bot_id] 

167 return True 

168 return False 

169 

170 async def deactivate(self, bot_id: str) -> bool: 

171 """Soft delete (set inactive). 

172 

173 Args: 

174 bot_id: Bot identifier 

175 

176 Returns: 

177 True if deactivated, False if not found 

178 """ 

179 async with self._lock: 

180 if bot_id in self._registrations: 

181 reg = self._registrations[bot_id] 

182 self._registrations[bot_id] = Registration( 

183 bot_id=reg.bot_id, 

184 config=reg.config, 

185 status="inactive", 

186 created_at=reg.created_at, 

187 updated_at=datetime.now(timezone.utc), 

188 last_accessed_at=reg.last_accessed_at, 

189 ) 

190 return True 

191 return False 

192 

193 async def list_active(self) -> list[Registration]: 

194 """List active registrations. 

195 

196 Returns: 

197 List of active Registration objects 

198 """ 

199 async with self._lock: 

200 return [ 

201 reg for reg in self._registrations.values() if reg.status == "active" 

202 ] 

203 

204 async def list_all(self) -> list[Registration]: 

205 """List all registrations. 

206 

207 Returns: 

208 List of all Registration objects 

209 """ 

210 async with self._lock: 

211 return list(self._registrations.values()) 

212 

213 async def list_ids(self) -> list[str]: 

214 """List active bot IDs. 

215 

216 Returns: 

217 List of active bot IDs 

218 """ 

219 async with self._lock: 

220 return [ 

221 reg.bot_id 

222 for reg in self._registrations.values() 

223 if reg.status == "active" 

224 ] 

225 

226 async def count(self) -> int: 

227 """Count active registrations. 

228 

229 Returns: 

230 Number of active registrations 

231 """ 

232 async with self._lock: 

233 return sum( 

234 1 for reg in self._registrations.values() if reg.status == "active" 

235 ) 

236 

237 async def clear(self) -> None: 

238 """Clear all registrations.""" 

239 async with self._lock: 

240 self._registrations.clear() 

241 

242 def __repr__(self) -> str: 

243 """String representation.""" 

244 return f"InMemoryBackend(count={len(self._registrations)})"