chore: add labeled client and update volume cleanup

framework/docker/chain_node_test.go

@@ -1,14 +1,14 @@
 package docker
 
 import (
+	tastoraclient "github.com/celestiaorg/tastora/framework/docker/client"
 	"github.com/stretchr/testify/require"
 	"testing"
 
 	"github.com/celestiaorg/tastora/framework/docker/container"
 	"github.com/celestiaorg/tastora/framework/docker/cosmos"
-	"go.uber.org/zap/zaptest"
 	"github.com/cosmos/cosmos-sdk/types/module/testutil"
-	dockerclient "github.com/moby/moby/client"
+	"go.uber.org/zap/zaptest"
 )
 
 func TestChainNodeHostName(t *testing.T) {
@@ -20,43 +20,43 @@ func TestChainNodeHostName(t *testing.T) {
 
 	// Create nodes with different indices
 	chainParams1 := cosmos.ChainNodeParams{
-		Validator:       true,
-		ChainID:         chainID,
-		BinaryName:      "test-binary",
-		CoinType:        "118",
-		GasPrices:       "0.025utia",
-		GasAdjustment:   1.0,
-		Env:             []string{},
+		Validator:           true,
+		ChainID:             chainID,
+		BinaryName:          "test-binary",
+		CoinType:            "118",
+		GasPrices:           "0.025utia",
+		GasAdjustment:       1.0,
+		Env:                 []string{},
 		AdditionalStartArgs: []string{},
-		EncodingConfig:  &testutil.TestEncodingConfig{},
+		EncodingConfig:      &testutil.TestEncodingConfig{},
 	}
-	node1 := cosmos.NewChainNode(logger, &dockerclient.Client{}, "test-network", testName, container.Image{}, "/test/home", 0, chainParams1)
+	node1 := cosmos.NewChainNode(logger, &tastoraclient.Client{}, "test-network", testName, container.Image{}, "/test/home", 0, chainParams1)
 
 	chainParams2 := cosmos.ChainNodeParams{
-		Validator:       true,
-		ChainID:         chainID,
-		BinaryName:      "test-binary",
-		CoinType:        "118",
-		GasPrices:       "0.025utia",
-		GasAdjustment:   1.0,
-		Env:             []string{},
+		Validator:           true,
+		ChainID:             chainID,
+		BinaryName:          "test-binary",
+		CoinType:            "118",
+		GasPrices:           "0.025utia",
+		GasAdjustment:       1.0,
+		Env:                 []string{},
 		AdditionalStartArgs: []string{},
-		EncodingConfig:  &testutil.TestEncodingConfig{},
+		EncodingConfig:      &testutil.TestEncodingConfig{},
 	}
-	node2 := cosmos.NewChainNode(logger, &dockerclient.Client{}, "test-network", testName, container.Image{}, "/test/home", 1, chainParams2)
+	node2 := cosmos.NewChainNode(logger, &tastoraclient.Client{}, "test-network", testName, container.Image{}, "/test/home", 1, chainParams2)
 
 	chainParams3 := cosmos.ChainNodeParams{
-		Validator:       false,
-		ChainID:         chainID,
-		BinaryName:      "test-binary",
-		CoinType:        "118",
-		GasPrices:       "0.025utia",
-		GasAdjustment:   1.0,
-		Env:             []string{},
+		Validator:           false,
+		ChainID:             chainID,
+		BinaryName:          "test-binary",
+		CoinType:            "118",
+		GasPrices:           "0.025utia",
+		GasAdjustment:       1.0,
+		Env:                 []string{},
 		AdditionalStartArgs: []string{},
-		EncodingConfig:  &testutil.TestEncodingConfig{},
+		EncodingConfig:      &testutil.TestEncodingConfig{},
 	}
-	node3 := cosmos.NewChainNode(logger, &dockerclient.Client{}, "test-network", testName, container.Image{}, "/test/home", 2, chainParams3)
+	node3 := cosmos.NewChainNode(logger, &tastoraclient.Client{}, "test-network", testName, container.Image{}, "/test/home", 2, chainParams3)
 
 	// get hostnames
 	hostname1 := node1.HostName()

framework/docker/client/client.go

@@ -0,0 +1,32 @@
+package client
+
+import (
+	"github.com/celestiaorg/tastora/framework/types"
+	"github.com/moby/moby/client"
+)
+
+// Client wraps a Docker client with an associated cleanup label.
+// the cleanup label is used to tag all resources (containers, volumes, networks)
+// created by this client, enabling cleanup to find and remove exactly the resources
+// associated with a specific test run.
+//
+// Client implements types.TastoraDockerClient.
+type Client struct {
+	*client.Client
+	cleanupLabel string
+}
+
+var _ types.TastoraDockerClient = (*Client)(nil)
+
+// NewClient creates a new Client with the given Docker client and cleanup label.
+func NewClient(c *client.Client, cleanupLabel string) *Client {
+	return &Client{
+		Client:       c,
+		cleanupLabel: cleanupLabel,
+	}
+}
+
+// CleanupLabel returns the cleanup label associated with this client.
+func (c *Client) CleanupLabel() string {
+	return c.cleanupLabel
+}

framework/docker/container/image.go

@@ -3,8 +3,8 @@ package container
 import (
 	"context"
 	"fmt"
+	"github.com/celestiaorg/tastora/framework/types"
 	dockerimagetypes "github.com/docker/docker/api/types/image"
-	"github.com/moby/moby/client"
 	"io"
 )
 
@@ -31,7 +31,7 @@ func (i Image) Ref() string {
 	return i.Repository + ":" + i.Version
 }
 
-func (i Image) PullImage(ctx context.Context, client *client.Client) error {
+func (i Image) PullImage(ctx context.Context, client types.TastoraDockerClient) error {
 	ref := i.Ref()
 	_, _, err := client.ImageInspectWithRaw(ctx, ref)
 	if err != nil {
@@ -43,4 +43,4 @@ func (i Image) PullImage(ctx context.Context, client *client.Client) error {
 		_ = rc.Close()
 	}
 	return nil
-}
+}

framework/docker/container/job.go

@@ -8,6 +8,7 @@ import (
 	"github.com/celestiaorg/tastora/framework/docker/consts"
 	"github.com/celestiaorg/tastora/framework/docker/internal"
 	"github.com/celestiaorg/tastora/framework/testutil/random"
+	"github.com/celestiaorg/tastora/framework/types"
 	"io"
 	"strconv"
 	"strings"
@@ -18,7 +19,6 @@ import (
 	dockerimagetypes "github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/api/types/network"
-	"github.com/moby/moby/client"
 	"github.com/moby/moby/errdefs"
 	"github.com/moby/moby/pkg/stdcopy"
 	"go.uber.org/zap"
@@ -27,7 +27,7 @@ import (
 // Job is a docker job runner.
 type Job struct {
 	log             *zap.Logger
-	client          *client.Client
+	client          types.TastoraDockerClient
 	repository, tag string
 
 	networkID string
@@ -42,7 +42,7 @@ type Job struct {
 // Most arguments (except tag) must be non-zero values or this function panics.
 // If tag is absent, defaults to "latest".
 // Currently, only public docker images are supported.
-func NewJob(logger *zap.Logger, cli *client.Client, networkID string, testName string, repository, tag string) *Job {
+func NewJob(logger *zap.Logger, cli types.TastoraDockerClient, networkID string, testName string, repository, tag string) *Job {
 	if logger == nil {
 		panic(errors.New("nil Logger"))
 	}
@@ -177,7 +177,7 @@ func (job *Job) CreateContainer(ctx context.Context, containerName, hostName str
 			Hostname: hostName,
 			User:     opts.User,
 
-			Labels: map[string]string{consts.CleanupLabel: job.testName},
+			Labels: map[string]string{consts.CleanupLabel: job.client.CleanupLabel()},
 		},
 		&container.HostConfig{
 			Binds:           opts.Binds,

framework/docker/container/lifecycle.go

@@ -19,7 +19,6 @@ import (
 	"github.com/docker/docker/api/types/network"
 	"github.com/docker/docker/api/types/volume"
 	"github.com/docker/go-connections/nat"
-	dockerclient "github.com/moby/moby/client"
 	"github.com/moby/moby/errdefs"
 	"go.uber.org/zap"
 )
@@ -29,16 +28,16 @@ var panicRe = regexp.MustCompile(`panic:.*\n`)
 
 type Lifecycle struct {
 	log               *zap.Logger
-	client            *dockerclient.Client
+	client            types.TastoraDockerClient
 	containerName     string
 	id                string
 	preStartListeners port.Listeners
 }
 
-func NewLifecycle(log *zap.Logger, client *dockerclient.Client, containerName string) *Lifecycle {
+func NewLifecycle(log *zap.Logger, c types.TastoraDockerClient, containerName string) *Lifecycle {
 	return &Lifecycle{
 		log:           log,
-		client:        client,
+		client:        c,
 		containerName: containerName,
 	}
 }
@@ -101,7 +100,7 @@ func (c *Lifecycle) CreateContainer(
 			Cmd:          cmd,
 			Env:          env,
 			Hostname:     hostName,
-			Labels:       map[string]string{consts.CleanupLabel: testName},
+			Labels:       map[string]string{consts.CleanupLabel: c.client.CleanupLabel()},
 			ExposedPorts: pS,
 		},
 		&container.HostConfig{
@@ -237,8 +236,8 @@ func (c *Lifecycle) RemoveContainer(ctx context.Context, opts ...types.RemoveOpt
 	return nil
 }
 
-func (c *Lifecycle) RemoveVolumes(ctx context.Context, cleanupLabel string) error {
-	filterArgs := filters.NewArgs(filters.Arg("label", fmt.Sprintf("%s=%s", consts.CleanupLabel, cleanupLabel)))
+func (c *Lifecycle) RemoveVolumes(ctx context.Context) error {
+	filterArgs := filters.NewArgs(filters.Arg("label", fmt.Sprintf("%s=%s", consts.CleanupLabel, c.client.CleanupLabel())))
 	volumeList, err := c.client.VolumeList(ctx, volume.ListOptions{Filters: filterArgs})
 	if err != nil {
 		return fmt.Errorf("failed to list volumes: %w", err)

framework/docker/container/node.go

@@ -12,15 +12,14 @@ import (
 	"github.com/docker/docker/api/types/mount"
 	volumetypes "github.com/docker/docker/api/types/volume"
 	"github.com/docker/go-connections/nat"
-	dockerclient "github.com/moby/moby/client"
 	"go.uber.org/zap"
 )
 
 // Node contains the fields and shared methods for docker nodes. (app nodes & bridge nodes)
 type Node struct {
 	VolumeName         string
 	NetworkID          string
-	DockerClient       *dockerclient.Client
+	DockerClient       types.TastoraDockerClient
 	TestName           string
 	Image              Image
 	ContainerLifecycle *Lifecycle
@@ -33,7 +32,7 @@ type Node struct {
 // NewNode creates a new Node instance with the required parameters.
 func NewNode(
 	networkID string,
-	dockerClient *dockerclient.Client,
+	dockerClient types.TastoraDockerClient,
 	testName string,
 	image Image,
 	homeDir string,
@@ -97,7 +96,7 @@ func (n *Node) RemoveContainer(ctx context.Context, opts ...types.RemoveOption)
 	if !removeOpts.RemoveVolumes {
 		return nil
 	}
-	return n.ContainerLifecycle.RemoveVolumes(ctx, n.TestName)
+	return n.ContainerLifecycle.RemoveVolumes(ctx)
 }
 
 // StopContainer gracefully stops the container associated with the Node using the provided context.
@@ -198,7 +197,7 @@ func (n *Node) ensureVolume(ctx context.Context, nodeName, volName string) error
 		v, createErr := n.DockerClient.VolumeCreate(ctx, volumetypes.CreateOptions{
 			Name: volName,
 			Labels: map[string]string{
-				consts.CleanupLabel:   n.TestName,
+				consts.CleanupLabel:   n.DockerClient.CleanupLabel(),
 				consts.NodeOwnerLabel: nodeName,
 			},
 		})

framework/docker/cosmos/chain_builder.go

@@ -14,7 +14,6 @@ import (
 	cryptocodec "github.com/cosmos/cosmos-sdk/crypto/codec"
 	"github.com/cosmos/cosmos-sdk/crypto/keyring"
 	"github.com/cosmos/cosmos-sdk/types/module/testutil"
-	"github.com/moby/moby/client"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest"
 )
@@ -115,7 +114,7 @@ type ChainBuilder struct {
 	// nodes is the array of node configurations that define the chain topology and individual node settings
 	nodes []ChainNodeConfig
 	// dockerClient is the Docker client instance used for all container operations (create, start, stop, etc.)
-	dockerClient *client.Client
+	dockerClient types.TastoraDockerClient
 	// dockerNetworkID is the ID of the Docker network where all chain nodes are deployed
 	dockerNetworkID string
 	// genesisBz contains raw bytes that should be written as the config/genesis.json file for the chain (optional)
@@ -247,8 +246,8 @@ func (b *ChainBuilder) WithNodes(nodeConfigs ...ChainNodeConfig) *ChainBuilder {
 }
 
 // WithDockerClient sets the Docker client
-func (b *ChainBuilder) WithDockerClient(client *client.Client) *ChainBuilder {
-	b.dockerClient = client
+func (b *ChainBuilder) WithDockerClient(c types.TastoraDockerClient) *ChainBuilder {
+	b.dockerClient = c
 	return b
 }
 

framework/docker/cosmos/chain_config.go

@@ -3,18 +3,18 @@ package cosmos
 import (
 	"context"
 	"github.com/celestiaorg/tastora/framework/docker/container"
+	"github.com/celestiaorg/tastora/framework/types"
 	"github.com/cosmos/cosmos-sdk/types/module/testutil"
-	"github.com/moby/moby/client"
 	"go.uber.org/zap"
 )
 
 type ChainConfig struct {
 	Logger *zap.Logger
 	// DockerClient is a Docker client instance used for the tests.
-	DockerClient *client.Client
+	DockerClient types.TastoraDockerClient
 	// DockerNetworkID is the ID of the docker network the nodes are deployed to.
 	DockerNetworkID string
-	
+
 	// Chain configuration fields (previously in ChainConfig)
 	// Chain name, e.g. celestia.
 	Name string

framework/docker/cosmos/node.go

@@ -32,7 +32,6 @@ import (
 	govv1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/docker/go-connections/nat"
-	dockerclient "github.com/moby/moby/client"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -116,7 +115,7 @@ type ChainNodeParams struct {
 // NewChainNode creates a new ChainNode with injected dependencies
 func NewChainNode(
 	logger *zap.Logger,
-	dockerClient *dockerclient.Client,
+	dockerClient types.TastoraDockerClient,
 	dockerNetworkID string,
 	testName string,
 	image container.Image,

framework/docker/dataavailability/config.go

@@ -2,7 +2,7 @@ package dataavailability
 
 import (
 	"github.com/celestiaorg/tastora/framework/docker/container"
-	"github.com/moby/moby/client"
+	"github.com/celestiaorg/tastora/framework/types"
 	"go.uber.org/zap"
 )
 
@@ -11,7 +11,7 @@ type Config struct {
 	// Logger is the logger instance used for all operations
 	Logger *zap.Logger
 	// DockerClient is the docker client instance
-	DockerClient *client.Client
+	DockerClient types.TastoraDockerClient
 	// DockerNetworkID is the ID of the docker network to use
 	DockerNetworkID string
 	// ChainID, e.g. test-chain

framework/docker/dataavailability/network_builder.go

@@ -3,10 +3,10 @@ package dataavailability
 import (
 	"context"
 	"fmt"
+	"github.com/celestiaorg/tastora/framework/types"
 	"testing"
 
 	"github.com/celestiaorg/tastora/framework/docker/container"
-	"github.com/moby/moby/client"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest"
 )
@@ -20,7 +20,7 @@ type NetworkBuilder struct {
 	// nodes is the array of node configurations that define the network topology and individual node settings
 	nodes []NodeConfig
 	// dockerClient is the Docker client instance used for all container operations
-	dockerClient *client.Client
+	dockerClient types.TastoraDockerClient
 	// dockerNetworkID is the ID of the Docker network where all Nodes are deployed
 	dockerNetworkID string
 	// logger is the structured logger for network operations and debugging. Defaults to test logger.
@@ -97,8 +97,8 @@ func (b *NetworkBuilder) WithBinaryName(binaryName string) *NetworkBuilder {
 }
 
 // WithDockerClient sets the Docker client
-func (b *NetworkBuilder) WithDockerClient(client *client.Client) *NetworkBuilder {
-	b.dockerClient = client
+func (b *NetworkBuilder) WithDockerClient(c types.TastoraDockerClient) *NetworkBuilder {
+	b.dockerClient = c
 	return b
 }