1+ package org .unicitylabs .sdk ;
2+
3+ import com .fasterxml .jackson .core .JsonProcessingException ;
4+ import okhttp3 .mockwebserver .Dispatcher ;
5+ import okhttp3 .mockwebserver .MockResponse ;
6+ import okhttp3 .mockwebserver .MockWebServer ;
7+ import okhttp3 .mockwebserver .RecordedRequest ;
8+ import com .fasterxml .jackson .databind .ObjectMapper ;
9+ import com .fasterxml .jackson .databind .JsonNode ;
10+ import org .jetbrains .annotations .Nullable ;
11+
12+ import java .io .IOException ;
13+ import java .util .Set ;
14+ import java .util .HashSet ;
15+ import java .util .UUID ;
16+
17+ public class MockAggregatorServer {
18+
19+ private final MockWebServer server ;
20+ private final ObjectMapper objectMapper ;
21+ private final Set <String > protectedMethods ;
22+ private volatile boolean simulateRateLimit = false ;
23+ private volatile int rateLimitRetryAfter = 0 ;
24+ private volatile String expectedApiKey = null ;
25+
26+ public MockAggregatorServer () {
27+ this .server = new MockWebServer ();
28+ this .objectMapper = new ObjectMapper ();
29+ this .protectedMethods = new HashSet <>();
30+ this .protectedMethods .add ("submit_commitment" );
31+
32+ server .setDispatcher (new Dispatcher () {
33+ @ Override
34+ public MockResponse dispatch (RecordedRequest request ) {
35+ return handleRequest (request );
36+ }
37+ });
38+ }
39+
40+ public void start () throws IOException {
41+ server .start ();
42+ }
43+
44+ public void shutdown () throws IOException {
45+ server .shutdown ();
46+ }
47+
48+ public String getUrl () {
49+ return server .url ("/" ).toString ();
50+ }
51+
52+ public RecordedRequest takeRequest () throws InterruptedException {
53+ return server .takeRequest ();
54+ }
55+
56+ public void simulateRateLimitForNextRequest (int retryAfterSeconds ) {
57+ this .simulateRateLimit = true ;
58+ this .rateLimitRetryAfter = retryAfterSeconds ;
59+ }
60+
61+ public void setExpectedApiKey (String apiKey ) {
62+ this .expectedApiKey = apiKey ;
63+ }
64+
65+ private MockResponse handleRequest (RecordedRequest request ) {
66+ try {
67+ if (simulateRateLimit ) {
68+ try {
69+ return new MockResponse ()
70+ .setResponseCode (429 )
71+ .setHeader ("Retry-After" , String .valueOf (rateLimitRetryAfter ))
72+ .setBody ("Too Many Requests" );
73+ } finally {
74+ // Reset for next request
75+ simulateRateLimit = false ;
76+ rateLimitRetryAfter = 0 ;
77+ }
78+ }
79+
80+ String method = extractJsonRpcMethod (request );
81+
82+ if (protectedMethods .contains (method ) && expectedApiKey != null && !hasValidApiKey (request )) {
83+ return new MockResponse ()
84+ .setResponseCode (401 )
85+ .setHeader ("WWW-Authenticate" , "Bearer" )
86+ .setBody ("Unauthorized" );
87+ }
88+
89+ return generateSuccessResponse (method );
90+
91+ } catch (Exception e ) {
92+ return new MockResponse ()
93+ .setResponseCode (400 )
94+ .setBody ("Bad Request" );
95+ }
96+ }
97+
98+ private boolean hasValidApiKey (RecordedRequest request ) {
99+ String authHeader = request .getHeader ("Authorization" );
100+ if (authHeader != null && authHeader .startsWith ("Bearer " )) {
101+ String providedKey = authHeader .substring (7 );
102+ return expectedApiKey .equals (providedKey );
103+ }
104+ return false ;
105+ }
106+
107+ private @ Nullable String extractJsonRpcMethod (RecordedRequest request ) throws JsonProcessingException {
108+ if (!"POST" .equals (request .getMethod ())) {
109+ return null ;
110+ }
111+ JsonNode jsonRequest = objectMapper .readTree (request .getBody ().readUtf8 ());
112+ return jsonRequest .has ("method" ) ? jsonRequest .get ("method" ).asText () : null ;
113+ }
114+
115+ private MockResponse generateSuccessResponse (String method ) {
116+ String responseBody ;
117+ String id = UUID .randomUUID ().toString ();
118+
119+ switch (method != null ? method : "" ) {
120+ case "submit_commitment" :
121+ responseBody = String .format (
122+ "{\n " +
123+ " \" jsonrpc\" : \" 2.0\" ,\n " +
124+ " \" result\" : {\n " +
125+ " \" status\" : \" SUCCESS\" \n " +
126+ " },\n " +
127+ " \" id\" : \" %s\" \n " +
128+ "}" , id );
129+ break ;
130+
131+ case "get_block_height" :
132+ responseBody = String .format (
133+ "{\n " +
134+ " \" jsonrpc\" : \" 2.0\" ,\n " +
135+ " \" result\" : {\n " +
136+ " \" blockNumber\" : \" 67890\" \n " +
137+ " },\n " +
138+ " \" id\" : \" %s\" \n " +
139+ "}" , id );
140+ break ;
141+
142+ default :
143+ responseBody = String .format (
144+ "{\n " +
145+ " \" jsonrpc\" : \" 2.0\" ,\n " +
146+ " \" result\" : \" OK\" ,\n " +
147+ " \" id\" : \" %s\" \n " +
148+ "}" , id );
149+ break ;
150+ }
151+
152+ return new MockResponse ()
153+ .setResponseCode (200 )
154+ .setHeader ("Content-Type" , "application/json" )
155+ .setBody (responseBody );
156+ }
157+ }
0 commit comments