|
7 | 7 | "net"
|
8 | 8 | "net/netip"
|
9 | 9 | "slices"
|
| 10 | + "strings" |
10 | 11 | "sync"
|
11 | 12 | "sync/atomic"
|
12 | 13 | "testing"
|
@@ -1451,10 +1452,35 @@ func (f *fakeDNSSetter) SetDNSHosts(hosts map[dnsname.FQDN][]netip.Addr) error {
|
1451 | 1452 | }
|
1452 | 1453 | }
|
1453 | 1454 |
|
| 1455 | +func newFakeUpdateHandler(ctx context.Context, t testing.TB) *fakeUpdateHandler { |
| 1456 | + return &fakeUpdateHandler{ |
| 1457 | + ctx: ctx, |
| 1458 | + t: t, |
| 1459 | + ch: make(chan *proto.WorkspaceUpdate), |
| 1460 | + } |
| 1461 | +} |
| 1462 | + |
| 1463 | +type fakeUpdateHandler struct { |
| 1464 | + ctx context.Context |
| 1465 | + t testing.TB |
| 1466 | + ch chan *proto.WorkspaceUpdate |
| 1467 | +} |
| 1468 | + |
| 1469 | +func (f *fakeUpdateHandler) Update(wu *proto.WorkspaceUpdate) error { |
| 1470 | + f.t.Helper() |
| 1471 | + select { |
| 1472 | + case <-f.ctx.Done(): |
| 1473 | + return timeoutOnFakeErr |
| 1474 | + case f.ch <- wu: |
| 1475 | + // OK |
| 1476 | + } |
| 1477 | + return nil |
| 1478 | +} |
| 1479 | + |
1454 | 1480 | func setupConnectedAllWorkspaceUpdatesController(
|
1455 | 1481 | ctx context.Context, t testing.TB, logger slog.Logger, opts ...tailnet.TunnelAllOption,
|
1456 | 1482 | ) (
|
1457 |
| - *fakeCoordinatorClient, *fakeWorkspaceUpdateClient, |
| 1483 | + *fakeCoordinatorClient, *fakeWorkspaceUpdateClient, *tailnet.TunnelAllWorkspaceUpdatesController, |
1458 | 1484 | ) {
|
1459 | 1485 | fConn := &fakeCoordinatee{}
|
1460 | 1486 | tsc := tailnet.NewTunnelSrcCoordController(logger, fConn)
|
@@ -1484,17 +1510,20 @@ func setupConnectedAllWorkspaceUpdatesController(
|
1484 | 1510 | err := testutil.RequireRecvCtx(ctx, t, updateCW.Wait())
|
1485 | 1511 | require.ErrorIs(t, err, io.EOF)
|
1486 | 1512 | })
|
1487 |
| - return coordC, updateC |
| 1513 | + return coordC, updateC, uut |
1488 | 1514 | }
|
1489 | 1515 |
|
1490 | 1516 | func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
|
1491 | 1517 | t.Parallel()
|
1492 | 1518 | ctx := testutil.Context(t, testutil.WaitShort)
|
1493 | 1519 | logger := testutil.Logger(t)
|
1494 | 1520 |
|
| 1521 | + fUH := newFakeUpdateHandler(ctx, t) |
1495 | 1522 | fDNS := newFakeDNSSetter(ctx, t)
|
1496 |
| - coordC, updateC := setupConnectedAllWorkspaceUpdatesController(ctx, t, logger, |
1497 |
| - tailnet.WithDNS(fDNS, "testy")) |
| 1523 | + coordC, updateC, updateCtrl := setupConnectedAllWorkspaceUpdatesController(ctx, t, logger, |
| 1524 | + tailnet.WithDNS(fDNS, "testy"), |
| 1525 | + tailnet.WithHandler(fUH), |
| 1526 | + ) |
1498 | 1527 |
|
1499 | 1528 | // Initial update contains 2 workspaces with 1 & 2 agents, respectively
|
1500 | 1529 | w1ID := testUUID(1)
|
@@ -1541,16 +1570,43 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
|
1541 | 1570 | dnsCall := testutil.RequireRecvCtx(ctx, t, fDNS.calls)
|
1542 | 1571 | require.Equal(t, expectedDNS, dnsCall.hosts)
|
1543 | 1572 | testutil.RequireSendCtx(ctx, t, dnsCall.err, nil)
|
| 1573 | + |
| 1574 | + // And the callback |
| 1575 | + cbUpdate := testutil.RequireRecvCtx(ctx, t, fUH.ch) |
| 1576 | + require.Equal(t, initUp, cbUpdate) |
| 1577 | + |
| 1578 | + // Current state should match |
| 1579 | + state := updateCtrl.CurrentState() |
| 1580 | + slices.SortFunc(state.UpsertedWorkspaces, func(a, b *proto.Workspace) int { |
| 1581 | + return strings.Compare(a.Name, b.Name) |
| 1582 | + }) |
| 1583 | + slices.SortFunc(state.UpsertedAgents, func(a, b *proto.Agent) int { |
| 1584 | + return strings.Compare(a.Name, b.Name) |
| 1585 | + }) |
| 1586 | + require.Equal(t, &proto.WorkspaceUpdate{ |
| 1587 | + UpsertedWorkspaces: []*proto.Workspace{ |
| 1588 | + {Id: w1ID[:], Name: "w1"}, |
| 1589 | + {Id: w2ID[:], Name: "w2"}, |
| 1590 | + }, |
| 1591 | + UpsertedAgents: []*proto.Agent{ |
| 1592 | + {Id: w1a1ID[:], Name: "w1a1", WorkspaceId: w1ID[:]}, |
| 1593 | + {Id: w2a1ID[:], Name: "w2a1", WorkspaceId: w2ID[:]}, |
| 1594 | + {Id: w2a2ID[:], Name: "w2a2", WorkspaceId: w2ID[:]}, |
| 1595 | + }, |
| 1596 | + }, state) |
1544 | 1597 | }
|
1545 | 1598 |
|
1546 | 1599 | func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
|
1547 | 1600 | t.Parallel()
|
1548 | 1601 | ctx := testutil.Context(t, testutil.WaitShort)
|
1549 | 1602 | logger := testutil.Logger(t)
|
1550 | 1603 |
|
| 1604 | + fUH := newFakeUpdateHandler(ctx, t) |
1551 | 1605 | fDNS := newFakeDNSSetter(ctx, t)
|
1552 |
| - coordC, updateC := setupConnectedAllWorkspaceUpdatesController(ctx, t, logger, |
1553 |
| - tailnet.WithDNS(fDNS, "testy")) |
| 1606 | + coordC, updateC, updateCtrl := setupConnectedAllWorkspaceUpdatesController(ctx, t, logger, |
| 1607 | + tailnet.WithDNS(fDNS, "testy"), |
| 1608 | + tailnet.WithHandler(fUH), |
| 1609 | + ) |
1554 | 1610 |
|
1555 | 1611 | w1ID := testUUID(1)
|
1556 | 1612 | w1a1ID := testUUID(1, 1)
|
@@ -1582,6 +1638,20 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
|
1582 | 1638 | require.Equal(t, expectedDNS, dnsCall.hosts)
|
1583 | 1639 | testutil.RequireSendCtx(ctx, t, dnsCall.err, nil)
|
1584 | 1640 |
|
| 1641 | + cbUpdate := testutil.RequireRecvCtx(ctx, t, fUH.ch) |
| 1642 | + require.Equal(t, initUp, cbUpdate) |
| 1643 | + |
| 1644 | + // Current state should match initial |
| 1645 | + state := updateCtrl.CurrentState() |
| 1646 | + require.Equal(t, &proto.WorkspaceUpdate{ |
| 1647 | + UpsertedWorkspaces: []*proto.Workspace{ |
| 1648 | + {Id: w1ID[:], Name: "w1"}, |
| 1649 | + }, |
| 1650 | + UpsertedAgents: []*proto.Agent{ |
| 1651 | + {Id: w1a1ID[:], Name: "w1a1", WorkspaceId: w1ID[:]}, |
| 1652 | + }, |
| 1653 | + }, state) |
| 1654 | + |
1585 | 1655 | // Send update that removes w1a1 and adds w1a2
|
1586 | 1656 | agentUpdate := &proto.WorkspaceUpdate{
|
1587 | 1657 | UpsertedAgents: []*proto.Agent{
|
@@ -1613,6 +1683,19 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
|
1613 | 1683 | dnsCall = testutil.RequireRecvCtx(ctx, t, fDNS.calls)
|
1614 | 1684 | require.Equal(t, expectedDNS, dnsCall.hosts)
|
1615 | 1685 | testutil.RequireSendCtx(ctx, t, dnsCall.err, nil)
|
| 1686 | + |
| 1687 | + cbUpdate = testutil.RequireRecvCtx(ctx, t, fUH.ch) |
| 1688 | + require.Equal(t, agentUpdate, cbUpdate) |
| 1689 | + |
| 1690 | + state = updateCtrl.CurrentState() |
| 1691 | + require.Equal(t, &proto.WorkspaceUpdate{ |
| 1692 | + UpsertedWorkspaces: []*proto.Workspace{ |
| 1693 | + {Id: w1ID[:], Name: "w1"}, |
| 1694 | + }, |
| 1695 | + UpsertedAgents: []*proto.Agent{ |
| 1696 | + {Id: w1a2ID[:], Name: "w1a2", WorkspaceId: w1ID[:]}, |
| 1697 | + }, |
| 1698 | + }, state) |
1616 | 1699 | }
|
1617 | 1700 |
|
1618 | 1701 | func TestTunnelAllWorkspaceUpdatesController_DNSError(t *testing.T) {
|
|
0 commit comments