test(#3996): cover mid-write cleanup, file modes, and missing parents

partialFailReader writes a genuine partial chunk on its first Read then
fails on every subsequent Read, so the mid-write-failure tests reach
atomicfile.Write's cleanup path AFTER a real os.CreateTemp'd file already
has partial content on disk, rather than relying on a permission trick
that prevents the temp file from ever being created. Also cover preserved
file modes and creation under missing parent directories.

This coverage backs the workspace media writer's reliance on
atomicfile.Write for never-overwrite, no-residue semantics.

Author: aheritier

pkg/atomicfile/atomicfile_test.go

@@ -1,18 +1,22 @@
-package atomicfile_test
+package atomicfile
 
 import (
 	"bytes"
+	"errors"
 	"os"
 	"path/filepath"
 	"runtime"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-
-	"github.com/docker/docker-agent/pkg/atomicfile"
 )
 
+// TestWriteCreatesFileWithMode is the plan's restored "new file gets the
+// requested mode" regression, inadvertently dropped when this file was
+// rewritten to add the mid-write cleanup tests below: Write must create a
+// brand-new file with exactly the requested permission bits, not whatever
+// natefinch/atomic.WriteFile's own default (umask-derived) mode would be.
 func TestWriteCreatesFileWithMode(t *testing.T) {
 	t.Parallel()
 	if runtime.GOOS == "windows" {
@@ -22,7 +26,7 @@ func TestWriteCreatesFileWithMode(t *testing.T) {
 	dir := t.TempDir()
 	path := filepath.Join(dir, "secret")
 
-	require.NoError(t, atomicfile.Write(path, bytes.NewReader([]byte("hello")), 0o600))
+	require.NoError(t, Write(path, bytes.NewReader([]byte("hello")), 0o600))
 
 	data, err := os.ReadFile(path)
 	require.NoError(t, err)
@@ -33,6 +37,10 @@ func TestWriteCreatesFileWithMode(t *testing.T) {
 	assert.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 }
 
+// TestWriteOverwritesAndRetightensMode is the restored counterpart for an
+// EXISTING destination with a looser mode (0o644): a replacement Write
+// must retighten it to exactly the requested (0o600) bits, not merely
+// preserve whatever mode the file already had.
 func TestWriteOverwritesAndRetightensMode(t *testing.T) {
 	t.Parallel()
 	if runtime.GOOS == "windows" {
@@ -43,7 +51,7 @@ func TestWriteOverwritesAndRetightensMode(t *testing.T) {
 	path := filepath.Join(dir, "secret")
 
 	require.NoError(t, os.WriteFile(path, []byte("old"), 0o644))
-	require.NoError(t, atomicfile.Write(path, bytes.NewReader([]byte("new")), 0o600))
+	require.NoError(t, Write(path, bytes.NewReader([]byte("new")), 0o600))
 
 	data, err := os.ReadFile(path)
 	require.NoError(t, err)
@@ -54,11 +62,84 @@ func TestWriteOverwritesAndRetightensMode(t *testing.T) {
 	assert.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 }
 
+// TestWriteReturnsErrorForMissingDirectory is the restored regression for
+// a missing parent directory: Write must surface the underlying error
+// rather than creating the parent itself or panicking.
 func TestWriteReturnsErrorForMissingDirectory(t *testing.T) {
 	t.Parallel()
 	dir := t.TempDir()
 	path := filepath.Join(dir, "missing", "file")
 
-	err := atomicfile.Write(path, bytes.NewReader([]byte("x")), 0o600)
+	err := Write(path, bytes.NewReader([]byte("x")), 0o600)
 	assert.Error(t, err)
 }
+
+// partialFailReader returns chunk on its first Read (so the underlying
+// os.CreateTemp'd file genuinely receives that data before anything goes
+// wrong) and then fails on every subsequent Read. This reproduces a
+// mid-write failure AFTER a real, partially written temporary file
+// already exists on disk — the seam the plan requires ("deterministically
+// fail only after a controlled temporary/partial output has been
+// created"), as opposed to a permission trick that would prevent the
+// temp file from ever being created at all.
+type partialFailReader struct {
+	chunk []byte
+	err   error
+	sent  bool
+}
+
+func (r *partialFailReader) Read(p []byte) (int, error) {
+	if !r.sent {
+		r.sent = true
+		n := copy(p, r.chunk)
+		return n, nil
+	}
+	return 0, r.err
+}
+
+// TestWrite_CleansUpTempFileOnMidWriteFailure proves the real production
+// cleanup path: after Write's underlying temp file has already received
+// partial data, a subsequent read failure must leave no temp file behind
+// in the target directory, and the destination path itself must remain
+// untouched (never partially written, never created if it didn't already
+// exist).
+func TestWrite_CleansUpTempFileOnMidWriteFailure(t *testing.T) {
+	dir := t.TempDir()
+	target := filepath.Join(dir, "artifact.bin")
+	injectedErr := errors.New("injected mid-write failure")
+
+	err := Write(target, &partialFailReader{chunk: []byte("partial-data"), err: injectedErr}, 0o600)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), injectedErr.Error())
+
+	_, statErr := os.Stat(target)
+	assert.True(t, os.IsNotExist(statErr), "the destination file must never be created on a mid-write failure")
+
+	entries, readErr := os.ReadDir(dir)
+	require.NoError(t, readErr)
+	assert.Empty(t, entries, "no partial/orphan temp file must remain in the directory after a mid-write failure")
+}
+
+// TestWrite_CleansUpTempFileOnMidWriteFailure_ExistingDestination is the
+// same regression when the destination already has prior content: a
+// mid-write failure for a REPLACEMENT write must leave the original
+// content intact and, again, no temp file residue.
+func TestWrite_CleansUpTempFileOnMidWriteFailure_ExistingDestination(t *testing.T) {
+	dir := t.TempDir()
+	target := filepath.Join(dir, "artifact.bin")
+	require.NoError(t, os.WriteFile(target, []byte("original"), 0o600))
+	injectedErr := errors.New("injected mid-write failure")
+
+	err := Write(target, &partialFailReader{chunk: []byte("partial-data"), err: injectedErr}, 0o600)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), injectedErr.Error())
+
+	data, readErr := os.ReadFile(target)
+	require.NoError(t, readErr)
+	assert.Equal(t, "original", string(data), "the original destination content must survive a failed replacement")
+
+	entries, err := os.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, entries, 1, "no partial/orphan temp file must remain alongside the untouched destination")
+	assert.Equal(t, "artifact.bin", entries[0].Name())
+}