Feat: add some improvements (#9)

* add truncate size option

* enable directory

* fix default truncate unit

* rename arg NewFileTruncator

Makefile

@@ -13,10 +13,11 @@ dist: $@
 
 .PHONY: clean test build
 
-test:
-	@mkdir -p test /tmp/test_results
+test: clean
+	@mkdir -p test /tmp/test_results cli/test/foo
+	touch cli/test/file0 cli/test/foo/file1 cli/test/foo/file2
 	gotestsum --junitfile /tmp/test_results/unit-tests.xml -- -coverprofile=./test/coverage.out ./...
 	go tool cover -html=test/coverage.out -o test/coverage.html
 
 clean:
-	- $(RM) -rf dist/* test/*
+	- $(RM) -rf dist/* test/* cli/test

cli/app.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path/filepath"
 	"time"
 
 	"github.com/elastic-infra/go-remove-slowly/filesystem"
@@ -21,6 +22,55 @@ func NewMyApp() *MyApp {
 	return &MyApp{app, nil}
 }
 
+func isDirectory(path string) (bool, error) {
+	info, err := os.Stat(path)
+	if err != nil {
+		return false, err
+	}
+	return info.IsDir(), nil
+}
+
+func getFilePaths(paths []string) ([]string, error) {
+	var files []string
+	checked := map[string]bool{}
+
+	walkFunc := func(path string, info os.FileInfo, err error) error {
+		isDir, err := isDirectory(path)
+		if err != nil {
+			return err
+		}
+		if isDir {
+			return nil
+		}
+		if _, ok := checked[path]; ok {
+			return nil
+		}
+		files = append(files, path)
+		checked[path] = true
+		return nil
+	}
+
+	for _, path := range paths {
+		isDir, err := isDirectory(path)
+		if err != nil {
+			return nil, err
+		}
+
+		if !isDir {
+			if _, ok := checked[path]; !ok {
+				files = append(files, path)
+				checked[path] = true
+			}
+			continue
+		}
+
+		if err := filepath.Walk(path, walkFunc); err != nil {
+			return nil, err
+		}
+	}
+	return files, nil
+}
+
 // NewApp returns a cli app
 func NewApp() *MyApp {
 	app := NewMyApp()
@@ -36,11 +86,17 @@ func NewApp() *MyApp {
 		stream := NewIoMayDumbWriter(os.Stdout, isDumb)
 		app.stream = stream
 		errs := []error{}
-		for _, filePath := range context.Args().Slice() {
-			fmt.Fprintln(app.stream, "Removing File: "+filePath)
-			truncator := filesystem.NewFileTruncator(filePath, context.Duration("interval"), app.stream)
+
+		targetFilePaths, err := getFilePaths(context.Args().Slice())
+		if err != nil {
+			return err
+		}
+
+		for _, target := range targetFilePaths {
+			fmt.Fprintln(app.stream, "Removing File: "+target)
+			truncator := filesystem.NewFileTruncator(target, context.Duration("interval"), context.Int64("size"), app.stream)
 			if err := truncator.Remove(); err != nil {
-				fmt.Fprintf(os.Stderr, "File %s removal error: %s\n", filePath, err.Error())
+				fmt.Fprintf(os.Stderr, "File %s removal error: %s\n", target, err.Error())
 				errs = append(errs, err)
 			}
 		}
@@ -71,6 +127,12 @@ func NewApp() *MyApp {
 			Aliases: []string{"v"},
 			Usage:   "Show version and build information",
 		},
+		&cli.Int64Flag{
+			Name:    "size",
+			Aliases: []string{"s"},
+			Usage:   "Truncation size at once[MB]",
+			Value:   filesystem.DefaultTruncateSizeMB,
+		},
 	}
 	return app
 }

cli/app_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"os"
+	"reflect"
 	"strings"
 	"testing"
 )
@@ -101,3 +102,131 @@ func createFile(path string, size int64) {
 		panic("Failed to write file " + path + " " + err.Error())
 	}
 }
+
+func Test_getFilePaths(t *testing.T) {
+	type args struct {
+		paths []string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    []string
+		wantErr bool
+	}{
+		{
+			name: "file only",
+			args: args{
+				paths: []string{"test/file0"},
+			},
+			want: []string{
+				"test/file0",
+			},
+		},
+		{
+			name: "dup files",
+			args: args{
+				paths: []string{"test/file0", "test/file0"},
+			},
+			want: []string{
+				"test/file0",
+			},
+		},
+		{
+			name: "dir only",
+			args: args{
+				paths: []string{"test/foo/"},
+			},
+			want: []string{
+				"test/foo/file1",
+				"test/foo/file2",
+			},
+		},
+		{
+			name: "dup dir",
+			args: args{
+				paths: []string{"test/", "test/foo/"},
+			},
+			want: []string{
+				"test/file0",
+				"test/foo/file1",
+				"test/foo/file2",
+			},
+		},
+		{
+			name: "file and dir",
+			args: args{
+				paths: []string{"test/file0", "test/foo/"},
+			},
+			want: []string{
+				"test/file0",
+				"test/foo/file1",
+				"test/foo/file2",
+			},
+		},
+		{
+			name: "no such file or dir",
+			args: args{
+				paths: []string{"test/xxxx"},
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := getFilePaths(tt.args.paths)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getFilePaths() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("getFilePaths() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func Test_isDirectory(t *testing.T) {
+	type args struct {
+		path string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    bool
+		wantErr bool
+	}{
+		{
+			name: "yes",
+			args: args{
+				path: "test",
+			},
+			want: true,
+		},
+		{
+			name: "not",
+			args: args{
+				path: "main.go",
+			},
+			want: false,
+		},
+		{
+			name: "error",
+			args: args{
+				path: "no_path",
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := isDirectory(tt.args.path)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("isDirectory() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("isDirectory() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

filesystem/truncate.go

@@ -9,6 +9,12 @@ import (
 	pb "gopkg.in/cheggaaa/pb.v1"
 )
 
+const (
+	// DefaultTruncateSizeMB represents default truncate size.
+	DefaultTruncateSizeMB = 1
+	truncateSizeUnit      = 1024 * 1024 // MB
+)
+
 // FileTruncator encapsulates necessary data for truncation
 type FileTruncator struct {
 	FilePath         string
@@ -19,10 +25,10 @@ type FileTruncator struct {
 }
 
 // NewFileTruncator returns a new file truncator
-func NewFileTruncator(filePath string, interval time.Duration, writer io.Writer) *FileTruncator {
+func NewFileTruncator(filePath string, interval time.Duration, sizeMB int64, writer io.Writer) *FileTruncator {
 	truncator := &FileTruncator{
 		FilePath:         filePath,
-		TruncateUnit:     1024 * 1024,
+		TruncateUnit:     sizeMB * truncateSizeUnit,
 		TruncateInterval: interval,
 		writer:           writer,
 	}

filesystem/truncate_test.go

@@ -8,21 +8,21 @@ import (
 )
 
 func TestNewFileTruncator(t *testing.T) {
-	truncator := NewFileTruncator("path", time.Duration(10), nil)
+	truncator := NewFileTruncator("path", time.Duration(10), DefaultTruncateSizeMB, nil)
 	if truncator.FilePath != "path" {
 		t.Fatalf("FilePath is incorrect")
 	}
 }
 
 func TestTruncateCount(t *testing.T) {
-	truncator := NewFileTruncator("path", time.Duration(10), nil)
+	truncator := NewFileTruncator("path", time.Duration(10), DefaultTruncateSizeMB, nil)
 	tests := []struct {
 		size  int64
 		count int
 	}{
-		{5 * 1024 * 1024, 5},
-		{5*1024*1024 + 5, 6},
-		{5*1024*1024 - 5, 5},
+		{5 * truncateSizeUnit, 5},
+		{5*truncateSizeUnit + 5, 6},
+		{5*truncateSizeUnit - 5, 5},
 		{512, 1},
 	}
 	for _, test := range tests {
@@ -43,12 +43,12 @@ func TestUpdateStat(t *testing.T) {
 		t.Fatalf("Failed to create file %s", err.Error())
 	}
 	var size int64
-	size = 10 * 1024 * 1024
+	size = 10 * truncateSizeUnit
 	_, err = file.WriteAt([]byte("a"), size-1)
 	if err != nil {
 		t.Fatalf("Failed to write to the file %s", err.Error())
 	}
-	truncator := NewFileTruncator(path, time.Duration(1), nil)
+	truncator := NewFileTruncator(path, time.Duration(1), DefaultTruncateSizeMB, nil)
 	err = truncator.UpdateStat()
 	if err != nil {
 		t.Fatalf("UpdateStat failed: %s", err.Error())
@@ -60,7 +60,7 @@ func TestUpdateStat(t *testing.T) {
 
 func TestUpdateStat_PathError(t *testing.T) {
 	path := fmt.Sprintf("%s/%s", os.TempDir(), "statTestFile")
-	truncator := NewFileTruncator(path, time.Duration(1), nil)
+	truncator := NewFileTruncator(path, time.Duration(1), DefaultTruncateSizeMB, nil)
 	err := truncator.UpdateStat()
 	if err == nil {
 		t.Fatal("Error did not happen")
@@ -75,12 +75,12 @@ func TestRemove(t *testing.T) {
 		t.Fatalf("Failed to create file %s", err.Error())
 	}
 	var size int64
-	size = 10 * 1024 * 1024
+	size = 10 * truncateSizeUnit
 	_, err = file.WriteAt([]byte("a"), size-1)
 	if err != nil {
 		t.Fatalf("Failed to write to the file %s", err.Error())
 	}
-	truncator := NewFileTruncator(path, time.Duration(1), nil)
+	truncator := NewFileTruncator(path, time.Duration(1), DefaultTruncateSizeMB, nil)
 	err = truncator.Remove()
 	if err != nil {
 		t.Fatalf("File Removal failed %s", err.Error())