"""Tests for the resource graph, admission gate and reachability prover. Run against your own work: pytest Run against the solution: LAB_MODULE=solution pytest """ from __future__ import annotations import importlib import os from typing import Any, Mapping, Optional import pytest lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab")) GraphError = lab.GraphError Resource = lab.Resource ResourceGraph = lab.ResourceGraph Action = lab.Action Change = lab.Change Plan = lab.Plan State = lab.State plan = lab.plan apply = lab.apply FORCE_NEW = lab.FORCE_NEW diff_attributes = lab.diff_attributes detect_drift = lab.detect_drift Effect = lab.Effect Policy = lab.Policy AdmissionGate = lab.AdmissionGate DEFAULT_POLICIES = lab.DEFAULT_POLICIES requires_private_endpoint = lab.requires_private_endpoint requires_signed_image = lab.requires_signed_image requires_digest_pin = lab.requires_digest_pin VNet = lab.VNet Subnet = lab.Subnet Peering = lab.Peering Direction = lab.Direction NsgRule = lab.NsgRule PrivateEndpoint = lab.PrivateEndpoint Service = lab.Service Workload = lab.Workload EgressPolicy = lab.EgressPolicy Topology = lab.Topology ReachabilityProver = lab.ReachabilityProver NodePool = lab.NodePool Workload3D = lab.Workload3D MIG_PROFILES = lab.MIG_PROFILES schedule = lab.schedule scale_out_delay = lab.scale_out_delay # ====================================================================================== # helpers # ====================================================================================== def ok_provider(address: str, action: Action, attributes: Optional[Mapping[str, Any]]) -> Mapping[str, Any]: return dict(attributes or {}) def failing_provider(*failing: str): def provider(address: str, action: Action, attributes: Optional[Mapping[str, Any]]) -> Mapping[str, Any]: if address in failing: raise RuntimeError(f"boom: {address}") return dict(attributes or {}) return provider def chain_graph() -> ResourceGraph: g = ResourceGraph() g.add(Resource("t.a", "t", {"v": 1})) g.add(Resource("t.b", "t", {"v": 2}, depends_on=("t.a",))) g.add(Resource("t.c", "t", {"v": 3}, depends_on=("t.b",))) return g def applied_state(graph: ResourceGraph) -> State: state = State() apply(plan(graph, state), graph, state, ok_provider) return state # ====================================================================================== # 1. the resource graph # ====================================================================================== def test_a_resource_can_be_added_and_read_back(): g = ResourceGraph() g.add(Resource("t.a", "t", {"v": 1})) assert g.get("t.a").attributes["v"] == 1 def test_a_duplicate_address_is_rejected(): g = ResourceGraph() g.add(Resource("t.a", "t")) with pytest.raises(GraphError, match="duplicate"): g.add(Resource("t.a", "t")) def test_an_unknown_address_raises(): with pytest.raises(GraphError): ResourceGraph().get("t.nope") def test_a_dependency_on_an_unknown_resource_is_rejected(): g = ResourceGraph() g.add(Resource("t.a", "t", depends_on=("t.ghost",))) with pytest.raises(GraphError, match="unknown"): g.validate() def test_a_cycle_is_rejected(): g = ResourceGraph() g.add(Resource("t.a", "t", depends_on=("t.b",))) g.add(Resource("t.b", "t", depends_on=("t.a",))) with pytest.raises(GraphError, match="cycle"): g.validate() def test_a_self_dependency_is_a_cycle(): g = ResourceGraph() g.add(Resource("t.a", "t", depends_on=("t.a",))) with pytest.raises(GraphError): g.validate() def test_a_longer_cycle_is_rejected(): g = ResourceGraph() g.add(Resource("t.a", "t", depends_on=("t.c",))) g.add(Resource("t.b", "t", depends_on=("t.a",))) g.add(Resource("t.c", "t", depends_on=("t.b",))) with pytest.raises(GraphError): g.validate() def test_dependencies_come_before_dependents(): order = chain_graph().order() assert order.index("t.a") < order.index("t.b") < order.index("t.c") def test_the_order_is_deterministic(): a, b = chain_graph(), chain_graph() assert a.order() == b.order() def test_independent_resources_are_ordered_by_address(): g = ResourceGraph() for name in ("t.z", "t.a", "t.m"): g.add(Resource(name, "t")) assert g.order() == ["t.a", "t.m", "t.z"] def test_transitive_dependents_are_the_blast_radius(): assert chain_graph().transitive_dependents("t.a") == ["t.b", "t.c"] def test_a_leaf_has_no_dependents(): assert chain_graph().transitive_dependents("t.c") == [] def test_a_diamond_reports_each_dependent_once(): g = ResourceGraph() g.add(Resource("t.a", "t")) g.add(Resource("t.b", "t", depends_on=("t.a",))) g.add(Resource("t.c", "t", depends_on=("t.a",))) g.add(Resource("t.d", "t", depends_on=("t.b", "t.c"))) assert g.transitive_dependents("t.a") == ["t.b", "t.c", "t.d"] # ====================================================================================== # 2. plan # ====================================================================================== def test_a_fresh_plan_creates_everything(): p = plan(chain_graph(), State()) assert len(p.by_action(Action.CREATE)) == 3 def test_a_plan_on_unchanged_state_is_empty(): g = chain_graph() assert plan(g, applied_state(g)).empty def test_an_attribute_change_is_an_update(): g = chain_graph() state = applied_state(g) g2 = chain_graph() g2._resources["t.b"] = Resource("t.b", "t", {"v": 99}, depends_on=("t.a",)) changes = plan(g2, state).by_action(Action.UPDATE) assert [c.address for c in changes] == ["t.b"] def test_an_update_names_the_changed_attributes(): g = chain_graph() state = applied_state(g) g2 = chain_graph() g2._resources["t.b"] = Resource("t.b", "t", {"v": 99, "w": 1}, depends_on=("t.a",)) change = plan(g2, state).by_action(Action.UPDATE)[0] assert change.changed_attributes == ("v", "w") def test_a_force_new_attribute_produces_a_replace(): g = ResourceGraph() g.add(Resource("azurerm_subnet.a", "azurerm_subnet", {"address_prefix": "10.0.0.0/24"})) state = applied_state(g) g2 = ResourceGraph() g2.add(Resource("azurerm_subnet.a", "azurerm_subnet", {"address_prefix": "10.0.1.0/24"})) change = plan(g2, state).changes[0] assert change.action is Action.REPLACE assert change.force_new_because == ("address_prefix",) def test_a_non_force_new_attribute_on_the_same_type_is_an_update(): g = ResourceGraph() g.add(Resource("azurerm_subnet.a", "azurerm_subnet", {"address_prefix": "10.0.0.0/24", "tags": {}})) state = applied_state(g) g2 = ResourceGraph() g2.add(Resource("azurerm_subnet.a", "azurerm_subnet", {"address_prefix": "10.0.0.0/24", "tags": {"a": "b"}})) assert plan(g2, state).changes[0].action is Action.UPDATE def test_a_replace_is_destructive(): g = ResourceGraph() g.add(Resource("azurerm_subnet.a", "azurerm_subnet", {"address_prefix": "10.0.0.0/24"})) state = applied_state(g) g2 = ResourceGraph() g2.add(Resource("azurerm_subnet.a", "azurerm_subnet", {"address_prefix": "10.0.9.0/24"})) assert len(plan(g2, state).destructive_changes) == 1 def test_a_resource_removed_from_the_config_is_destroyed(): g = chain_graph() state = applied_state(g) smaller = ResourceGraph() smaller.add(Resource("t.a", "t", {"v": 1})) deletes = plan(smaller, state).by_action(Action.DELETE) assert {c.address for c in deletes} == {"t.b", "t.c"} def test_deletes_are_ordered_dependents_first(): g = chain_graph() state = applied_state(g) smaller = ResourceGraph() smaller.add(Resource("t.a", "t", {"v": 1})) deletes = [c.address for c in plan(smaller, state).by_action(Action.DELETE)] assert deletes.index("t.c") < deletes.index("t.b") def test_a_plan_validates_the_graph(): g = ResourceGraph() g.add(Resource("t.a", "t", depends_on=("t.b",))) g.add(Resource("t.b", "t", depends_on=("t.a",))) with pytest.raises(GraphError): plan(g, State()) def test_the_summary_counts_every_action(): summary = plan(chain_graph(), State()).summary() assert "3 to add" in summary def test_diff_attributes_reports_added_removed_and_changed(): assert diff_attributes({"a": 1, "b": 2}, {"b": 3, "c": 4}) == ["a", "b", "c"] def test_diff_attributes_is_empty_for_identical_maps(): assert diff_attributes({"a": 1}, {"a": 1}) == [] # ====================================================================================== # 3. apply # ====================================================================================== def test_a_clean_apply_applies_everything(): g = chain_graph() result = apply(plan(g, State()), g, State(), ok_provider) assert result.ok and len(result.applied) == 3 def test_apply_writes_state(): g = chain_graph() state = State() apply(plan(g, state), g, state, ok_provider) assert state.addresses() == ["t.a", "t.b", "t.c"] def test_apply_follows_dependency_order(): g = chain_graph() result = apply(plan(g, State()), g, State(), ok_provider) assert list(result.applied) == ["t.a", "t.b", "t.c"] def test_a_failure_is_reported(): g = chain_graph() result = apply(plan(g, State()), g, State(), failing_provider("t.b")) assert not result.ok and result.failed == ("t.b",) def test_a_failure_skips_its_dependents(): g = chain_graph() result = apply(plan(g, State()), g, State(), failing_provider("t.b")) assert result.skipped == ("t.c",) def test_a_failure_does_not_skip_independent_resources(): g = ResourceGraph() g.add(Resource("t.a", "t")) g.add(Resource("t.b", "t")) result = apply(plan(g, State()), g, State(), failing_provider("t.a")) assert "t.b" in result.applied def test_a_skipped_resource_is_not_written_to_state(): g = chain_graph() state = State() apply(plan(g, state), g, state, failing_provider("t.b")) assert state.get("t.c") is None def test_a_successful_resource_before_the_failure_is_kept_in_state(): g = chain_graph() state = State() apply(plan(g, state), g, state, failing_provider("t.b")) assert state.get("t.a") is not None def test_the_error_is_recorded_per_address(): g = chain_graph() result = apply(plan(g, State()), g, State(), failing_provider("t.b")) assert "boom" in result.errors["t.b"] def test_a_transitive_failure_skips_the_whole_chain(): g = ResourceGraph() g.add(Resource("t.a", "t")) g.add(Resource("t.b", "t", depends_on=("t.a",))) g.add(Resource("t.c", "t", depends_on=("t.b",))) g.add(Resource("t.d", "t", depends_on=("t.c",))) result = apply(plan(g, State()), g, State(), failing_provider("t.a")) assert set(result.skipped) == {"t.b", "t.c", "t.d"} def test_deletes_remove_from_state(): g = chain_graph() state = applied_state(g) smaller = ResourceGraph() smaller.add(Resource("t.a", "t", {"v": 1})) apply(plan(smaller, state), smaller, state, ok_provider) assert state.addresses() == ["t.a"] # ====================================================================================== # 4. drift # ====================================================================================== def test_matching_state_and_reality_is_clean(): state = applied_state(chain_graph()) assert detect_drift(state, state.snapshot()).clean def test_a_changed_attribute_is_drift(): state = applied_state(chain_graph()) reality = state.snapshot() reality["t.b"] = {**reality["t.b"], "v": 99} report = detect_drift(state, reality) assert len(report.drifted) == 1 and report.drifted[0].attribute == "v" def test_drift_reports_both_values(): state = applied_state(chain_graph()) reality = state.snapshot() reality["t.b"] = {**reality["t.b"], "v": 99} d = detect_drift(state, reality).drifted[0] assert (d.in_state, d.in_reality) == (2, 99) def test_a_resource_gone_from_reality_is_missing(): state = applied_state(chain_graph()) reality = state.snapshot() del reality["t.c"] assert detect_drift(state, reality).missing == ("t.c",) def test_a_resource_not_in_state_is_unmanaged(): state = applied_state(chain_graph()) reality = state.snapshot() reality["t.rogue"] = {"v": 0} assert detect_drift(state, reality).unmanaged == ("t.rogue",) def test_all_three_categories_are_reported_together(): state = applied_state(chain_graph()) reality = state.snapshot() reality["t.b"] = {**reality["t.b"], "v": 99} del reality["t.c"] reality["t.rogue"] = {} report = detect_drift(state, reality) assert report.drifted and report.missing and report.unmanaged assert not report.clean def test_drift_is_reported_per_attribute(): state = applied_state(chain_graph()) reality = state.snapshot() reality["t.a"] = {"v": 9, "extra": 1} assert len(detect_drift(state, reality).drifted) == 2 def test_an_empty_state_reports_everything_as_unmanaged(): report = detect_drift(State(), {"t.a": {}, "t.b": {}}) assert report.unmanaged == ("t.a", "t.b") # ====================================================================================== # 5. the admission gate # ====================================================================================== def resource(**attributes) -> Resource: rtype = attributes.pop("resource_type", "kubernetes_deployment") address = attributes.pop("address", "r.x") return Resource(address, rtype, attributes) def deny_reasons(res: Resource, gate: AdmissionGate = None): gate = gate or AdmissionGate() return [d for d in gate.evaluate(res) if not d.allowed] def test_a_compliant_resource_is_allowed(): res = resource(image="reg/x@sha256:ab", image_signature="s", image_signer="platform-ci", cpu_limit="1", memory_limit="1Gi", tags={"owner": "layla"}) assert not deny_reasons(res) def test_restricted_data_without_a_private_endpoint_is_denied(): res = resource(resource_type="azurerm_storage_account", data_classification="restricted", private_endpoint=False, tags={"owner": "l"}) assert any(d.rule == "restricted-data-private-endpoint" for d in deny_reasons(res)) def test_a_private_endpoint_with_public_access_still_enabled_is_denied(): res = resource(resource_type="azurerm_key_vault", data_classification="restricted", private_endpoint=True, public_network_access="enabled", tags={"owner": "l"}) denials = deny_reasons(res) assert any("does not close the public path" in d.reason for d in denials) def test_a_private_endpoint_with_public_access_disabled_is_allowed(): res = resource(resource_type="azurerm_key_vault", data_classification="restricted", private_endpoint=True, public_network_access="disabled", tags={"owner": "l"}) assert not deny_reasons(res) def test_internal_data_needs_no_private_endpoint(): res = resource(resource_type="azurerm_storage_account", data_classification="internal", tags={"owner": "l"}) assert not deny_reasons(res) def test_an_unsigned_image_is_denied(): res = resource(image="reg/x@sha256:ab", cpu_limit="1", memory_limit="1Gi", tags={"owner": "l"}) assert any(d.rule == "signed-images" for d in deny_reasons(res)) def test_an_image_signed_by_an_untrusted_signer_is_denied(): res = resource(image="reg/x@sha256:ab", image_signature="s", image_signer="mallory", cpu_limit="1", memory_limit="1Gi", tags={"owner": "l"}) assert any("not trusted" in d.reason for d in deny_reasons(res)) def test_a_mutable_tag_is_denied(): res = resource(image="reg/x:latest", image_signature="s", image_signer="platform-ci", cpu_limit="1", memory_limit="1Gi", tags={"owner": "l"}) assert any(d.rule == "digest-pinned-images" for d in deny_reasons(res)) def test_missing_resource_limits_are_denied(): res = resource(image="reg/x@sha256:ab", image_signature="s", image_signer="platform-ci", tags={"owner": "l"}) assert any(d.rule == "resource-limits" for d in deny_reasons(res)) def test_a_missing_owner_tag_is_denied_on_any_type(): res = resource(resource_type="anything.at.all") assert any(d.rule == "owner-tag" for d in deny_reasons(res)) def test_a_public_ip_on_compute_is_denied(): res = resource(resource_type="azurerm_linux_virtual_machine", public_ip="20.1.2.3", tags={"owner": "l"}) assert any(d.rule == "no-public-ip" for d in deny_reasons(res)) def test_wildcard_egress_is_denied(): res = resource(resource_type="azurerm_firewall_policy_rule_collection", egress_destinations=("*",), tags={"owner": "l"}) assert any(d.rule == "no-wildcard-egress" for d in deny_reasons(res)) def test_every_violated_rule_is_reported_not_only_the_first(): res = resource(image="reg/x:latest") rules = {d.rule for d in deny_reasons(res)} assert len(rules) >= 4 def test_every_denial_names_its_rule_and_the_address(): res = resource(address="k8s.tool", image="reg/x:latest") for d in deny_reasons(res): assert d.rule and d.address == "k8s.tool" def test_a_policy_that_raises_denies_rather_than_failing_open(): def explodes(res): raise RuntimeError("policy bug") gate = AdmissionGate([Policy("boom", (), explodes)]) denials = deny_reasons(resource(), gate) assert denials and "policy evaluation failed" in denials[0].reason def test_a_policy_applies_only_to_its_declared_types(): p = Policy("only-storage", ("azurerm_storage_account",), lambda r: "no") gate = AdmissionGate([p]) assert not deny_reasons(resource(resource_type="kubernetes_deployment"), gate) assert deny_reasons(resource(resource_type="azurerm_storage_account"), gate) def test_an_empty_applies_to_matches_every_type(): gate = AdmissionGate([Policy("all", (), lambda r: "no")]) assert deny_reasons(resource(resource_type="anything"), gate) def test_admitting_a_graph_returns_every_denial(): g = ResourceGraph() g.add(Resource("a.x", "kubernetes_deployment", {"image": "x:latest"})) g.add(Resource("b.y", "azurerm_linux_virtual_machine", {"public_ip": "1.2.3.4"})) ok, denials = AdmissionGate().admit(g) assert not ok and len({d.address for d in denials}) == 2 def test_a_fully_compliant_graph_is_admitted(): g = ResourceGraph() g.add(Resource("a.x", "kubernetes_deployment", {"image": "reg/x@sha256:ab", "image_signature": "s", "image_signer": "platform-ci", "cpu_limit": "1", "memory_limit": "1Gi", "tags": {"owner": "l"}})) ok, denials = AdmissionGate().admit(g) assert ok and denials == [] # ====================================================================================== # 6. NSG evaluation # ====================================================================================== def nsg_topology() -> Topology: t = Topology() t.add_vnet(VNet("v", "uaenorth", "10.0.0.0/16")) t.add_subnet(Subnet("s", "v", "10.0.1.0/24")) return t def test_the_first_matching_rule_by_priority_wins(): t = nsg_topology() t.add_nsg_rule(NsgRule("deny", 100, Direction.OUTBOUND, Effect.DENY, "*", "*"), "s") t.add_nsg_rule(NsgRule("allow", 200, Direction.OUTBOUND, Effect.ALLOW, "*", "*"), "s") assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "x", 443) == (Effect.DENY, "deny") def test_swapping_the_priorities_swaps_the_outcome(): t = nsg_topology() t.add_nsg_rule(NsgRule("deny", 200, Direction.OUTBOUND, Effect.DENY, "*", "*"), "s") t.add_nsg_rule(NsgRule("allow", 100, Direction.OUTBOUND, Effect.ALLOW, "*", "*"), "s") assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "x", 443)[0] is Effect.ALLOW def test_no_matching_rule_is_an_implicit_deny(): effect, rule = nsg_topology().nsg_verdict("s", Direction.OUTBOUND, "s", "x", 443) assert effect is Effect.DENY and rule == "implicit-deny" def test_direction_is_respected(): t = nsg_topology() t.add_nsg_rule(NsgRule("in", 100, Direction.INBOUND, Effect.ALLOW, "*", "*"), "s") assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "x", 443)[0] is Effect.DENY def test_a_port_list_narrows_a_rule(): t = nsg_topology() t.add_nsg_rule(NsgRule("https", 100, Direction.OUTBOUND, Effect.ALLOW, "*", "*", (443,)), "s") assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "x", 443)[0] is Effect.ALLOW assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "x", 22)[0] is Effect.DENY def test_an_empty_port_list_matches_any_port(): t = nsg_topology() t.add_nsg_rule(NsgRule("any", 100, Direction.OUTBOUND, Effect.ALLOW, "*", "*"), "s") assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "x", 22)[0] is Effect.ALLOW def test_a_cidr_destination_matches_an_address_inside_it(): t = nsg_topology() t.add_nsg_rule(NsgRule("net", 100, Direction.OUTBOUND, Effect.ALLOW, "*", "10.0.2.0/24"), "s") assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "10.0.2.7", 443)[0] is Effect.ALLOW assert t.nsg_verdict("s", Direction.OUTBOUND, "s", "10.0.3.7", 443)[0] is Effect.DENY # ====================================================================================== # 7. topology construction # ====================================================================================== def test_a_subnet_referencing_an_unknown_vnet_is_rejected(): t = Topology() with pytest.raises(GraphError): t.add_subnet(Subnet("s", "ghost", "10.0.1.0/24")) def test_a_peering_referencing_an_unknown_vnet_is_rejected(): t = nsg_topology() with pytest.raises(GraphError): t.add_peering(Peering("v", "ghost")) def test_peering_is_directional(): t = nsg_topology() t.add_vnet(VNet("w", "uaenorth", "10.1.0.0/16")) t.add_peering(Peering("v", "w")) assert t.peered("v") == ["w"] and t.peered("w") == [] def test_peering_is_not_transitive(): t = nsg_topology() t.add_vnet(VNet("w", "uaenorth", "10.1.0.0/16")) t.add_vnet(VNet("x", "uaenorth", "10.2.0.0/16")) t.add_peering(Peering("v", "w")) t.add_peering(Peering("w", "x")) assert t.peered("v") == ["w"] def test_the_region_of_a_subnet_comes_from_its_vnet(): t = nsg_topology() assert t.region_of_subnet("s") == "uaenorth" # ====================================================================================== # 8. the reachability prover # ====================================================================================== def allow_all(t: Topology, *subnets: str) -> None: for s in subnets: t.add_nsg_rule(NsgRule("out", 100, Direction.OUTBOUND, Effect.ALLOW, "*", "*"), s) t.add_nsg_rule(NsgRule("in", 100, Direction.INBOUND, Effect.ALLOW, "*", "*"), s) def simple_topology() -> Topology: t = Topology() t.add_vnet(VNet("platform", "uaenorth", "10.10.0.0/16")) t.add_subnet(Subnet("aks", "platform", "10.10.1.0/24")) t.add_subnet(Subnet("endpoints", "platform", "10.10.2.0/24")) allow_all(t, "aks", "endpoints") t.add_service(Service("openai", "uaenorth", public_network_access=False)) t.add_workload(Workload("agent", "aks")) return t def test_a_private_endpoint_in_the_same_subnet_is_reachable(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe", "aks", "openai", "10.10.1.4")) assert ReachabilityProver(t).reach("agent", "openai").reachable def test_a_private_endpoint_in_another_subnet_of_the_same_vnet_is_reachable(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe", "endpoints", "openai", "10.10.2.4")) assert ReachabilityProver(t).reach("agent", "openai").reachable def test_a_service_with_no_endpoint_and_no_egress_is_unreachable(): assert not ReachabilityProver(simple_topology()).reach("agent", "openai").reachable def test_an_unreachable_result_explains_why(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe", "aks", "openai", "10.10.1.4")) t.nsg_rules["aks"] = [NsgRule("deny", 50, Direction.OUTBOUND, Effect.DENY, "*", "*")] result = ReachabilityProver(t).reach("agent", "openai") assert not result.reachable and result.blocked_by def test_a_reachable_result_carries_a_counter_example_path(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe", "aks", "openai", "10.10.1.4")) path = ReachabilityProver(t).reach("agent", "openai").counter_example assert path is not None and "private-endpoint:pe" in path.render() def test_a_private_endpoint_without_a_linked_dns_zone_does_not_count(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe", "aks", "openai", "10.10.1.4", dns_zone_linked=False)) result = ReachabilityProver(t).reach("agent", "openai") assert not result.reachable assert any("DNS" in r for r in result.blocked_by) def test_a_public_service_is_reachable_through_permitted_egress(): t = simple_topology() t.add_service(Service("public-api", "westeurope", public_network_access=True)) t.set_egress(EgressPolicy("aks", allow_internet=True)) assert ReachabilityProver(t).reach("agent", "public-api").reachable def test_egress_to_an_unlisted_host_is_blocked(): t = simple_topology() t.add_service(Service("public-api", "westeurope", public_network_access=True)) t.set_egress(EgressPolicy("aks", allowed_fqdns=frozenset({"something-else"}))) assert not ReachabilityProver(t).reach("agent", "public-api").reachable def test_an_allow_listed_host_is_reachable(): t = simple_topology() t.add_service(Service("docs", "westeurope", public_network_access=True)) t.set_egress(EgressPolicy("aks", allowed_fqdns=frozenset({"docs"}))) assert ReachabilityProver(t).reach("agent", "docs").reachable def test_a_service_with_public_access_disabled_is_not_reachable_via_egress(): t = simple_topology() t.set_egress(EgressPolicy("aks", allow_internet=True)) assert not ReachabilityProver(t).reach("agent", "openai").reachable def peered_topology() -> Topology: t = Topology() t.add_vnet(VNet("platform", "uaenorth", "10.10.0.0/16")) t.add_vnet(VNet("legacy", "westeurope", "10.30.0.0/16")) t.add_subnet(Subnet("aks", "platform", "10.10.1.0/24")) t.add_subnet(Subnet("legacy-app", "legacy", "10.30.1.0/24")) allow_all(t, "aks", "legacy-app") t.add_peering(Peering("platform", "legacy")) t.add_workload(Workload("agent", "aks")) t.add_service(Service("external", "westeurope", public_network_access=True)) t.set_egress(EgressPolicy("aks", allowed_fqdns=frozenset())) t.set_egress(EgressPolicy("legacy-app", allow_internet=True)) return t def test_a_path_through_a_peering_is_found(): result = ReachabilityProver(peered_topology()).reach("agent", "external") assert result.reachable def test_the_peering_path_is_visible_in_the_counter_example(): path = ReachabilityProver(peered_topology()).reach("agent", "external").counter_example assert "peering:platform->legacy" in path.render() def test_a_naive_per_subnet_check_would_miss_it(): t = peered_topology() assert not t.egress["aks"].allow_internet # the naive check says "safe" assert ReachabilityProver(t).reach("agent", "external").reachable def test_the_residency_proof_finds_the_offending_path(): result = ReachabilityProver(peered_topology()).prove_no_egress("agent", "external") assert result.reachable and result.counter_example.leaves_region def test_the_residency_proof_is_clean_when_only_a_private_endpoint_exists(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe", "aks", "openai", "10.10.1.4")) assert not ReachabilityProver(t).prove_no_egress("agent", "openai").reachable def test_an_nsg_deny_on_the_peering_hop_blocks_the_path(): t = peered_topology() t.nsg_rules["aks"] = [ NsgRule("deny-legacy", 100, Direction.OUTBOUND, Effect.DENY, "*", "legacy-app"), NsgRule("allow", 200, Direction.OUTBOUND, Effect.ALLOW, "*", "*")] assert not ReachabilityProver(t).reach("agent", "external").reachable def test_an_inbound_deny_at_the_far_end_blocks_the_path(): t = peered_topology() t.nsg_rules["legacy-app"] = [ NsgRule("deny-in", 100, Direction.INBOUND, Effect.DENY, "*", "*"), NsgRule("allow-out", 100, Direction.OUTBOUND, Effect.ALLOW, "*", "*")] assert not ReachabilityProver(t).reach("agent", "external").reachable def test_an_unknown_workload_raises(): with pytest.raises(GraphError): ReachabilityProver(simple_topology()).reach("ghost", "openai") def test_an_unknown_service_raises(): with pytest.raises(GraphError): ReachabilityProver(simple_topology()).reach("agent", "ghost") def test_the_prover_is_deterministic(): t = peered_topology() a = ReachabilityProver(t).reach("agent", "external") b = ReachabilityProver(t).reach("agent", "external") assert a.counter_example.render() == b.counter_example.render() def test_paths_are_returned_shortest_first(): t = simple_topology() t.add_private_endpoint(PrivateEndpoint("pe1", "aks", "openai", "10.10.1.4")) t.add_private_endpoint(PrivateEndpoint("pe2", "endpoints", "openai", "10.10.2.4")) paths = ReachabilityProver(t).reach("agent", "openai").paths assert len(paths[0].hops) <= len(paths[-1].hops) # ====================================================================================== # 9. GPU scheduling # ====================================================================================== def pool(**kw) -> NodePool: base = dict(name="p", gpu_type="A100-80", gpus_per_node=8, node_count=2, max_nodes=4) base.update(kw) return NodePool(**base) def test_a_pool_without_mig_has_one_slice_per_gpu(): assert pool().slices_per_gpu == 1 def test_mig_partitions_a_gpu(): assert pool(mig_profile="1g.10gb").slices_per_gpu == 7 def test_total_slices_multiplies_out(): assert pool(node_count=2, gpus_per_node=8, mig_profile="2g.20gb").total_slices == 48 def test_an_unknown_mig_profile_falls_back_to_one(): assert pool(mig_profile="nonsense").slices_per_gpu == 1 def test_every_mig_profile_divides_the_gpu(): for count, gib in MIG_PROFILES.values(): assert count * gib <= 80 def test_scale_out_delay_includes_engine_warmup(): p = pool(startup_seconds=300) assert scale_out_delay(p, engine_warmup_seconds=120) == 420 def test_a_workload_is_placed_on_a_matching_pool(): w = Workload3D("w", 2, "A100-80", toleration=("nvidia.com/gpu",)) placements = schedule([w], [pool(taints=("nvidia.com/gpu",))]) assert placements[0].placed def test_a_workload_without_the_toleration_is_not_placed(): w = Workload3D("w", 2, "A100-80") placements = schedule([w], [pool(taints=("nvidia.com/gpu",))]) assert not placements[0].placed and "taint" in placements[0].reason def test_a_workload_needing_a_different_gpu_type_is_not_placed(): w = Workload3D("w", 1, "H100-80") placements = schedule([w], [pool()]) assert not placements[0].placed and "H100-80" in placements[0].reason def test_a_gang_is_placed_across_whole_nodes(): w = Workload3D("tp8", 8, "A100-80", gang=True) placements = schedule([w], [pool(gpus_per_node=8, node_count=2)]) assert placements[0].placed and placements[0].nodes_used == 1 def test_a_gang_larger_than_one_node_spans_nodes(): w = Workload3D("tp16", 16, "A100-80", gang=True) placements = schedule([w], [pool(gpus_per_node=8, node_count=2)]) assert placements[0].placed and placements[0].nodes_used == 2 def test_a_gang_that_does_not_fit_is_not_partially_placed(): w = Workload3D("tp16", 16, "A100-80", gang=True) placements = schedule([w], [pool(gpus_per_node=8, node_count=1)]) assert not placements[0].placed def test_a_gang_cannot_use_fragmented_capacity(): small = Workload3D("a", 1, "A100-80") gang = Workload3D("g", 8, "A100-80", gang=True) placements = {p.workload: p for p in schedule([small, gang], [pool(node_count=1)])} assert placements["g"].placed # placed first, being larger assert not placements["a"].placed # nothing left def test_larger_workloads_are_placed_first(): small = Workload3D("small", 1, "A100-80") large = Workload3D("large", 8, "A100-80", gang=True) placements = schedule([small, large], [pool(node_count=1, gpus_per_node=8)]) assert {p.workload for p in placements if p.placed} == {"large"} def test_capacity_is_consumed_across_workloads(): ws = [Workload3D(f"w{i}", 4, "A100-80") for i in range(5)] placements = schedule(ws, [pool(node_count=2, gpus_per_node=8)]) assert sum(1 for p in placements if p.placed) == 4 def test_placements_are_returned_sorted_by_workload(): ws = [Workload3D("z", 1, "A100-80"), Workload3D("a", 1, "A100-80")] assert [p.workload for p in schedule(ws, [pool()])] == ["a", "z"] def test_scheduling_is_deterministic(): ws = [Workload3D(f"w{i}", 2, "A100-80") for i in range(6)] pools = [pool(name="p1"), pool(name="p2")] assert schedule(ws, pools) == schedule(ws, pools) def test_an_unplaced_workload_always_has_a_reason(): ws = [Workload3D("w", 99, "A100-80")] assert schedule(ws, [pool()])[0].reason