test(connlib): assign at least one gateway to each site (#6201)

I believe this is the cause of some flakiness in the proptests.
Sometimes, we would end up with a site that did not have any gateways
assigned.

Example of a failing test:
https://github.com/firezone/firezone/actions/runs/10280367072/job/28447560561?pr=6200.
This commit is contained in:
Thomas Eizinger
2024-08-09 03:16:04 +01:00
committed by GitHub
parent da296038b9
commit bc1ab58c93

View File

@@ -9,9 +9,8 @@ use connlib_shared::{
messages::{
client::{
ResourceDescriptionCidr, ResourceDescriptionDns, ResourceDescriptionInternet, Site,
SiteId,
},
GatewayId, RelayId,
RelayId,
},
proptest::{
any_ip_network, cidr_resource, dns_resource, domain_name, gateway_id, relay_id, site,
@@ -61,7 +60,6 @@ pub(crate) fn latency(max: u64) -> impl Strategy<Value = Duration> {
pub(crate) fn stub_portal() -> impl Strategy<Value = StubPortal> {
collection::hash_set(site(), 1..=3)
.prop_flat_map(|sites| {
let gateway_site = any_site(sites.clone()).prop_map(|s| s.id);
let cidr_resources = collection::hash_set(
cidr_resource_outside_reserved_ranges(any_site(sites.clone())),
1..5,
@@ -74,20 +72,14 @@ pub(crate) fn stub_portal() -> impl Strategy<Value = StubPortal> {
],
1..5,
);
let internet_resource = internet_resource(any_site(sites));
let internet_resource = internet_resource(any_site(sites.clone()));
// Gateways are unique across sites.
// Generate a map with `GatewayId`s as keys and then flip it into a map of site -> set(gateways).
let gateways_by_site = collection::hash_map(gateway_id(), gateway_site, 1..=3)
.prop_map(|gateway_site| {
let mut gateways_by_site = HashMap::<SiteId, HashSet<GatewayId>>::default();
for (gid, sid) in gateway_site {
gateways_by_site.entry(sid).or_default().insert(gid);
}
gateways_by_site
});
// Assign between 1 and 3 gateways to each site.
let gateways_by_site = sites
.into_iter()
.map(|site| (Just(site.id), collection::hash_set(gateway_id(), 1..=3)))
.collect::<Vec<_>>()
.prop_map(HashMap::from_iter);
let gateway_selector = any::<sample::Selector>();