Coverage for mcpgateway / toolops / utils / db_util.py: 100%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/toolops/utils/db_util.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Jay Bandlamudi 

6 

7ContextForge - Main module for handling toolops related database operations. 

8 

9This module defines the utility funtions to read/write/update toolops related database tables. 

10""" 

11 

12# Third-Party 

13from sqlalchemy.orm import Session 

14 

15# First-Party 

16from mcpgateway.db import Tool 

17from mcpgateway.db import ToolOpsTestCases as TestCaseRecord 

18from mcpgateway.services.logging_service import LoggingService 

19from mcpgateway.utils.services_auth import decode_auth 

20 

21logging_service = LoggingService() 

22logger = logging_service.get_logger(__name__) 

23 

24 

25def populate_testcases_table(tool_id, test_cases, run_status, db: Session): 

26 """ 

27 Method to write and update toolops test cases to database table 

28 

29 Args: 

30 tool_id: unqiue Tool ID used in MCP-CF 

31 test_cases: list of generated test cases, each test case is a dictionary object 

32 run_status: status of test case generation request such as in-progess, complete , failed 

33 db: DB session to access the database 

34 

35 Examples: 

36 >>> from unittest.mock import MagicMock, patch 

37 >>> # Setup: Get the current module path dynamically to ensure patch works regardless of file location 

38 >>> mod_path = populate_testcases_table.__module__ 

39 

40 >>> # Case 1: Insert New Record (Tool ID not found in DB) 

41 >>> mock_db = MagicMock() 

42 >>> # Simulate query returning None (record does not exist) 

43 >>> mock_db.query.return_value.filter_by.return_value.first.return_value = None 

44 

45 >>> # Patch the TestCaseRecord class specifically in THIS module 

46 >>> with patch(f"{mod_path}.TestCaseRecord") as MockRecord: 

47 ... populate_testcases_table("tool-123", [{"test": "case"}], "in-progress", mock_db) 

48 ... 

49 ... # Verify the class was instantiated 

50 ... MockRecord.assert_called_with(tool_id="tool-123", test_cases=[{"test": "case"}], run_status="in-progress") 

51 ... # Verify DB interactions 

52 ... mock_db.add.assert_called() 

53 ... mock_db.commit.assert_called() 

54 

55 >>> # Case 2: Update Existing Record 

56 >>> mock_db_update = MagicMock() 

57 >>> existing_record = MagicMock() 

58 >>> mock_db_update.query.return_value.filter_by.return_value.first.return_value = existing_record 

59 

60 >>> # We still patch TestCaseRecord to ensure no side effects, though it's not instantiated here 

61 >>> with patch(f"{mod_path}.TestCaseRecord"): 

62 ... populate_testcases_table("tool-123", [{"test": "new"}], "completed", mock_db_update) 

63 ... 

64 ... # Verify fields were updated on existing record 

65 ... assert existing_record.test_cases == [{"test": "new"}] 

66 ... assert existing_record.run_status == "completed" 

67 ... # Verify add was NOT called 

68 ... mock_db_update.add.assert_not_called() 

69 """ 

70 tool_record = db.query(TestCaseRecord).filter_by(tool_id=tool_id).first() 

71 if not tool_record: 

72 test_case_record = TestCaseRecord(tool_id=tool_id, test_cases=test_cases, run_status=run_status) 

73 # Add to DB 

74 db.add(test_case_record) 

75 db.commit() 

76 db.refresh(test_case_record) 

77 logger.info("Added tool test case record with empty test cases for tool " + str(tool_id) + " with status " + str(run_status)) 

78 # elif tool_record and test_cases != [] and run_status == 'completed': 

79 elif tool_record: 

80 tool_record.test_cases = test_cases 

81 tool_record.run_status = run_status 

82 db.commit() 

83 db.refresh(tool_record) 

84 logger.info("Updated tool record in table with test cases for tool " + str(tool_id) + " with status " + str(run_status)) 

85 

86 

87def query_testcases_table(tool_id, db: Session): 

88 """ 

89 Method to read toolops test cases from database table 

90 

91 Args: 

92 tool_id: unqiue Tool ID used in MCP-CF 

93 db: DB session to access the database 

94 

95 Returns: 

96 This method returns tool record for specified tool id and tool record contains 'tool_id','test_cases','run_status'. 

97 

98 Examples: 

99 >>> from unittest.mock import MagicMock, patch 

100 >>> mock_db = MagicMock() 

101 

102 >>> # Create a dummy record to return 

103 >>> mock_record = MagicMock() 

104 >>> mock_record.tool_id = "tool-abc" 

105 >>> mock_record.test_cases = [{"input": "test"}] 

106 

107 >>> # Mock the chain: db.query(...).filter_by(...).first() 

108 >>> mock_db.query.return_value.filter_by.return_value.first.return_value = mock_record 

109 

110 >>> # Execute 

111 >>> result = query_testcases_table("tool-abc", mock_db) 

112 

113 >>> # Verify result and calls 

114 >>> result.tool_id 

115 'tool-abc' 

116 >>> mock_db.query.assert_called() 

117 """ 

118 tool_record = db.query(TestCaseRecord).filter_by(tool_id=tool_id).first() 

119 logger.info("Tool record obtained from table for tool - " + str(tool_id)) 

120 return tool_record 

121 

122 

123def query_tool_auth(tool_id, db: Session): 

124 """ 

125 Method to read tools table from database and get tool auth 

126 

127 Args: 

128 tool_id: unique Tool ID used in MCP-CF 

129 db: DB session to access the database 

130 

131 Returns: 

132 This method returns tool auth specified tool id. 

133 

134 Examples: 

135 >>> from unittest.mock import MagicMock, patch 

136 >>> mod_path = query_tool_auth.__module__ 

137 

138 >>> # Case 1: Successful Auth Retrieval 

139 >>> mock_db = MagicMock() 

140 >>> mock_tool_record = MagicMock() 

141 >>> mock_tool_record.auth_value = "encoded-val" 

142 >>> mock_db.query.return_value.filter_by.return_value.first.return_value = mock_tool_record 

143 

144 >>> # We nest the patches to avoid SyntaxError in doctest multiline 'with' statements 

145 >>> with patch(f"{mod_path}.decode_auth", side_effect=lambda x: f"decoded-{x}"): 

146 ... with patch(f"{mod_path}.Tool"): 

147 ... auth = query_tool_auth("tool-1", mock_db) 

148 ... print(auth) 

149 decoded-encoded-val 

150 

151 >>> # Case 2: Exception Handling 

152 >>> import logging 

153 >>> logging.disable(logging.CRITICAL) 

154 >>> mock_db_fail = MagicMock() 

155 >>> mock_db_fail.query.side_effect = Exception("DB Connection Error") 

156 

157 >>> with patch(f"{mod_path}.Tool"): 

158 ... auth = query_tool_auth("tool-2", mock_db_fail) 

159 ... print(auth) 

160 None 

161 >>> logging.disable(logging.NOTSET) 

162 """ 

163 tool_auth = None 

164 try: 

165 tool_record = db.query(Tool).filter_by(id=tool_id).first() 

166 tool_auth = decode_auth(tool_record.auth_value) 

167 logger.info("Tool auth obtained from table for the tool - " + str(tool_id)) 

168 except Exception as e: 

169 logger.error("Error in obtaining authorization for the tool - " + tool_id + " , " + str(e)) 

170 return tool_auth 

171 

172 

173# if __name__=='__main__': 

174# # First-Party 

175# from mcpgateway.db import SessionLocal 

176# from mcpgateway.services.tool_service import ToolService 

177 

178# tool_id = '36451eb11de64ebf8f224fc41a846ff0' 

179# tool_service = ToolService() 

180# db = SessionLocal() 

181 

182# tool_auth = query_tool_auth(tool_id, db) 

183# print(tool_auth)