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

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

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

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Jay Bandlamudi 

6 

7MCP Gateway - 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# Third-Party 

12from sqlalchemy.orm import Session 

13 

14# First-Party 

15from mcpgateway.db import Tool 

16from mcpgateway.db import ToolOpsTestCases as TestCaseRecord 

17from mcpgateway.services.logging_service import LoggingService 

18from mcpgateway.utils.services_auth import decode_auth 

19 

20logging_service = LoggingService() 

21logger = logging_service.get_logger(__name__) 

22 

23 

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

25 """ 

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

27 

28 Args: 

29 tool_id: unqiue Tool ID used in MCP-CF 

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

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

32 db: DB session to access the database 

33 

34 Examples: 

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

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

37 >>> mod_path = populate_testcases_table.__module__ 

38 

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

40 >>> mock_db = MagicMock() 

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

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

43 

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

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

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

47 ... 

48 ... # Verify the class was instantiated 

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

50 ... # Verify DB interactions 

51 ... mock_db.add.assert_called() 

52 ... mock_db.commit.assert_called() 

53 

54 >>> # Case 2: Update Existing Record 

55 >>> mock_db_update = MagicMock() 

56 >>> existing_record = MagicMock() 

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

58 

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

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

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

62 ... 

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

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

65 ... assert existing_record.run_status == "completed" 

66 ... # Verify add was NOT called 

67 ... mock_db_update.add.assert_not_called() 

68 """ 

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

70 if not tool_record: 

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

72 # Add to DB 

73 db.add(test_case_record) 

74 db.commit() 

75 db.refresh(test_case_record) 

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

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

78 elif tool_record: 78 ↛ exitline 78 didn't return from function 'populate_testcases_table' because the condition on line 78 was always true

79 tool_record.test_cases = test_cases 

80 tool_record.run_status = run_status 

81 db.commit() 

82 db.refresh(tool_record) 

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

84 

85 

86def query_testcases_table(tool_id, db: Session): 

87 """ 

88 Method to read toolops test cases from database table 

89 

90 Args: 

91 tool_id: unqiue Tool ID used in MCP-CF 

92 db: DB session to access the database 

93 

94 Returns: 

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

96 

97 Examples: 

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

99 >>> mock_db = MagicMock() 

100 

101 >>> # Create a dummy record to return 

102 >>> mock_record = MagicMock() 

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

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

105 

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

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

108 

109 >>> # Execute 

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

111 

112 >>> # Verify result and calls 

113 >>> result.tool_id 

114 'tool-abc' 

115 >>> mock_db.query.assert_called() 

116 """ 

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

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

119 return tool_record 

120 

121 

122def query_tool_auth(tool_id, db: Session): 

123 """ 

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

125 

126 Args: 

127 tool_id: unique Tool ID used in MCP-CF 

128 db: DB session to access the database 

129 

130 Returns: 

131 This method returns tool auth specified tool id. 

132 

133 Examples: 

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

135 >>> mod_path = query_tool_auth.__module__ 

136 

137 >>> # Case 1: Successful Auth Retrieval 

138 >>> mock_db = MagicMock() 

139 >>> mock_tool_record = MagicMock() 

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

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

142 

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

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

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

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

147 ... print(auth) 

148 decoded-encoded-val 

149 

150 >>> # Case 2: Exception Handling 

151 >>> mock_db_fail = MagicMock() 

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

153 

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

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

156 ... print(auth) 

157 None 

158 """ 

159 tool_auth = None 

160 try: 

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

162 tool_auth = decode_auth(tool_record.auth_value) 

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

164 except Exception as e: 

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

166 return tool_auth 

167 

168 

169# if __name__=='__main__': 

170# # First-Party 

171# from mcpgateway.db import SessionLocal 

172# from mcpgateway.services.tool_service import ToolService 

173 

174# tool_id = '36451eb11de64ebf8f224fc41a846ff0' 

175# tool_service = ToolService() 

176# db = SessionLocal() 

177 

178# tool_auth = query_tool_auth(tool_id, db) 

179# print(tool_auth)