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

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

2

3 from __future__ import annotations

4

5 import asyncio

6 from datetime import datetime, timezone

7 from typing import Any

8

9 from .models import Registration

10

11

12 class 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)})"