1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
use std::collections::HashMap;
use std::time::Duration;

use async_compression::tokio::write::GzipEncoder;
use maplit::hashmap;
use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use tokio::io::AsyncWriteExt;
use url::Url;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;

use crate::uplink::Endpoints;
use crate::uplink::UplinkConfig;

/// Get a query ID, body, and a PQ manifest with that ID and body.
pub fn fake_manifest() -> (String, String, HashMap<String, String>) {
    let id = "1234".to_string();
    let body = r#"query { typename }"#.to_string();
    let manifest = hashmap! { id.to_string() => body.to_string() };
    (id, body, manifest)
}

/// Mocks an uplink server with a persisted query list containing no operations.
pub async fn mock_empty_pq_uplink() -> (UplinkMockGuard, UplinkConfig) {
    mock_pq_uplink(&HashMap::new()).await
}

/// Mocks an uplink server with a persisted query list with a delay.
pub async fn mock_pq_uplink_with_delay(
    manifest: &HashMap<String, String>,
    delay: Duration,
) -> (UplinkMockGuard, UplinkConfig) {
    let (guard, url) = mock_pq_uplink_one_endpoint(manifest, Some(delay)).await;
    (
        guard,
        UplinkConfig::for_tests(Endpoints::fallback(vec![url])),
    )
}

/// Mocks an uplink server with a persisted query list containing operations passed to this function.
pub async fn mock_pq_uplink(manifest: &HashMap<String, String>) -> (UplinkMockGuard, UplinkConfig) {
    let (guard, url) = mock_pq_uplink_one_endpoint(manifest, None).await;
    (
        guard,
        UplinkConfig::for_tests(Endpoints::fallback(vec![url])),
    )
}

/// Guards for the uplink and GCS mock servers, dropping these structs shuts down the server.
pub struct UplinkMockGuard {
    _uplink_mock_guard: MockServer,
    _gcs_mock_guard: MockServer,
}

#[derive(Deserialize, Serialize)]
struct Operation {
    id: String,
    body: String,
}

/// Mocks an uplink server; returns a single Url rather than a full UplinkConfig, so you
/// can combine it with another one to test failover.
pub async fn mock_pq_uplink_one_endpoint(
    manifest: &HashMap<String, String>,
    delay: Option<Duration>,
) -> (UplinkMockGuard, Url) {
    let operations: Vec<Operation> = manifest
        // clone the manifest so the caller can still make assertions about it
        .clone()
        .drain()
        .map(|(id, body)| Operation { id, body })
        .collect();

    let mock_gcs_server = MockServer::start().await;

    let body_json = serde_json::to_vec(&json!({
      "format": "apollo-persisted-query-manifest",
      "version": 1,
      "operations": operations
    }))
    .expect("Failed to convert into body.");
    let mut encoder = GzipEncoder::new(Vec::new());
    encoder.write_all(&body_json).await.unwrap();
    encoder.shutdown().await.unwrap();
    let compressed_body = encoder.into_inner();

    let gcs_response = ResponseTemplate::new(200)
        .set_body_raw(compressed_body, "application/octet-stream")
        .append_header("Content-Encoding", "gzip");

    Mock::given(method("GET"))
        .and(header("Accept-Encoding", "gzip"))
        .respond_with(gcs_response)
        .mount(&mock_gcs_server)
        .await;

    let mock_gcs_server_uri: Url = mock_gcs_server.uri().parse().unwrap();

    let mock_uplink_server = MockServer::start().await;

    let mut uplink_response = ResponseTemplate::new(200).set_body_json(json!({
          "data": {
            "persistedQueries": {
              "__typename": "PersistedQueriesResult",
              "id": "889406d7-b4f8-44df-a499-6c1e3c1bea09:1",
              "minDelaySeconds": 60,
              "chunks": [
                {
                  "id": "graph-id/889406a1-b4f8-44df-a499-6c1e3c1bea09/ec8ae3ae3eb00c738031dbe81603489b5d24fbf58f15bdeec1587282ee4e6eea",
                  "urls": [
                    "https://a.broken.gcs.url.that.will.get.fetched.and.skipped.unknown/",
                    mock_gcs_server_uri,
                  ]
                }
              ]
            }
          }
        }));

    if let Some(delay) = delay {
        uplink_response = uplink_response.set_delay(delay);
    }

    Mock::given(method("POST"))
        .respond_with(uplink_response)
        .mount(&mock_uplink_server)
        .await;

    let url = mock_uplink_server.uri().parse().unwrap();
    (
        UplinkMockGuard {
            _uplink_mock_guard: mock_uplink_server,
            _gcs_mock_guard: mock_gcs_server,
        },
        url,
    )
}

/// Mocks an uplink server which returns bad GCS URLs.
pub async fn mock_pq_uplink_bad_gcs() -> (MockServer, Url) {
    let mock_uplink_server = MockServer::start().await;

    let  uplink_response = ResponseTemplate::new(200).set_body_json(json!({
          "data": {
            "persistedQueries": {
              "__typename": "PersistedQueriesResult",
              "id": "889406d7-b4f8-44df-a499-6c1e3c1bea09:1",
              "minDelaySeconds": 60,
              "chunks": [
                {
                  "id": "graph-id/889406a1-b4f8-44df-a499-6c1e3c1bea09/ec8ae3ae3eb00c738031dbe81603489b5d24fbf58f15bdeec1587282ee4e6eea",
                  "urls": [
                    "https://definitely.not.gcs.unknown"
                  ]
                }
              ]
            }
          }
        }));

    Mock::given(method("POST"))
        .respond_with(uplink_response)
        .mount(&mock_uplink_server)
        .await;

    let url = mock_uplink_server.uri().parse().unwrap();
    (mock_uplink_server, url)
}