From 2d70a8ed31ef0718b65a79eed4cb7a5badeea455 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Tue, 18 Feb 2025 06:46:32 +1100
Subject: [PATCH] test(connlib): create dedicated Internet site (#8153)

To ensure that our test suite represents production as much as possible,
we introduce a dedicated "Internet" site into the `StubPortal` that only
hosts the Internet resource. All other creates resources are assigned to
other sites.
---
 .../tunnel/proptest-regressions/tests.txt     |  1 +
 rust/connlib/tunnel/src/tests/strategies.rs   | 27 +++++++++++++------
 rust/connlib/tunnel/src/tests/sut.rs          |  2 +-
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/rust/connlib/tunnel/proptest-regressions/tests.txt b/rust/connlib/tunnel/proptest-regressions/tests.txt
index 40dac8d46..5b94736bf 100644
--- a/rust/connlib/tunnel/proptest-regressions/tests.txt
+++ b/rust/connlib/tunnel/proptest-regressions/tests.txt
@@ -155,3 +155,4 @@ cc 893e70e809ad210d5ed7aab465bda2dbf2c8c472e6bdf63e66812fa238dd8966
 cc 3a21d576c6ab4baded47975d9e4acf91303c32a1d679ebdc6e6f5bc029712e3c
 cc 737a635b47b3b7fa128e9adbe5d45e18ea97b5633841b1616dedc0ccf8b61d16
 cc c29319940567b848032f49e1ce46100af145f028354403a0baf566d29ad20006
+cc a7f22e7cc2c79ffd580baf4bc8296557c67afe245ccf07e895e7cd2a969a228e
diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs
index 34bda4532..24043e0a6 100644
--- a/rust/connlib/tunnel/src/tests/strategies.rs
+++ b/rust/connlib/tunnel/src/tests/strategies.rs
@@ -11,6 +11,7 @@ use ip_network::{IpNetwork, Ipv4Network, Ipv6Network};
 use itertools::Itertools;
 use prop::sample;
 use proptest::{collection, prelude::*};
+use std::iter;
 use std::{
     collections::{BTreeMap, BTreeSet},
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
@@ -72,25 +73,27 @@ pub(crate) fn latency(max: u64) -> impl Strategy<Value = Duration> {
 /// Similar as in production, the portal holds a list of DNS and CIDR resources (those are also sampled from the given sites).
 /// Via this site mapping, these resources are implicitly assigned to a gateway.
 pub(crate) fn stub_portal() -> impl Strategy<Value = StubPortal> {
-    collection::btree_set(site(), 1..=3)
+    collection::btree_set(site(), 2..=4)
         .prop_flat_map(|sites| {
+            let (internet_site, regular_sites) = create_internet_site(sites);
+
             let cidr_resources = collection::btree_set(
-                cidr_resource_outside_reserved_ranges(any_site(sites.clone())),
+                cidr_resource_outside_reserved_ranges(any_site(regular_sites.clone())),
                 1..5,
             );
             let dns_resources = collection::btree_set(
                 prop_oneof![
-                    non_wildcard_dns_resource(any_site(sites.clone())),
-                    star_wildcard_dns_resource(any_site(sites.clone())),
-                    double_star_wildcard_dns_resource(any_site(sites.clone())),
+                    non_wildcard_dns_resource(any_site(regular_sites.clone())),
+                    star_wildcard_dns_resource(any_site(regular_sites.clone())),
+                    double_star_wildcard_dns_resource(any_site(regular_sites.clone())),
                 ],
                 1..5,
             );
-            let internet_resource = internet_resource(any_site(sites.clone()));
+            let internet_resource = internet_resource(Just(internet_site.clone()));
 
             // Assign between 1 and 3 gateways to each site.
-            let gateways_by_site = sites
-                .into_iter()
+            let gateways_by_site = iter::once(internet_site)
+                .chain(regular_sites)
                 .map(|site| (Just(site.id), collection::btree_set(gateway_id(), 1..=3)))
                 .collect::<Vec<_>>()
                 .prop_map(BTreeMap::from_iter);
@@ -124,6 +127,14 @@ pub(crate) fn stub_portal() -> impl Strategy<Value = StubPortal> {
         )
 }
 
+fn create_internet_site(mut sites: BTreeSet<Site>) -> (Site, BTreeSet<Site>) {
+    // Rebrand the first site as the Internet site. That way, we can guarantee to always have one.
+    let mut internet_site = sites.pop_first().unwrap();
+    internet_site.name = "Internet".to_owned();
+
+    (internet_site, sites)
+}
+
 pub(crate) fn relays(
     id: impl Strategy<Value = RelayId>,
 ) -> impl Strategy<Value = BTreeMap<RelayId, Host<u64>>> {
diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs
index 9170a88c3..1975dc1f4 100644
--- a/rust/connlib/tunnel/src/tests/sut.rs
+++ b/rust/connlib/tunnel/src/tests/sut.rs
@@ -296,7 +296,7 @@ impl TunnelTest {
                 let packets_per_sec = num_packets / num_seconds / num_connections;
 
                 // This has been chosen through experimentation. It primarily serves as a regression tool to ensure our idle-traffic doesn't suddenly spike.
-                const THRESHOLD: f64 = 2.1;
+                const THRESHOLD: f64 = 2.3;
 
                 if packets_per_sec > THRESHOLD {
                     tracing::error!("Expected at most {THRESHOLD} packets / sec in the network while idling. Got: {packets_per_sec}");