Testing cli tool with logging

In the previous article I described how to start and manage processes. The next step is to implement a cli utility that uses the resulting code. The task of implementing such a utility is not very difficult, but this is how to test it - this is already a more interesting. In this article I will describe one of the ways of testing that involves reading the output logs.

Flags

To start the necessary process it must be something passed to the utility. In addition, configuration values are required to start, for example

t Time to live of lock.
k Lock key name
l Logging level

and probably something else will be added in the future. The final kind of launch command would like to see such
djob [options] - command.
Command is the command (with arguments to run). -- should separate the arguments of the utility from the arguments of the program to run.
There are several excellent solutions for working with command-line arguments, such as spf13/cobra or urfave/cli. But they’re good for building an interface out of a lot of commands, and for one (as in my example) they’re redundant. So I used the flag library.
I have defined the following structure with arguments:

type CliParams struct {
    key      *string
    cmdArgs  []string
    logLevel *string
    ttl      *int
}

And made a function that parses arguments:

func parseCliParams(args []string) (*CliParams, *flag.FlagSet, error) {
    flags := flag.NewFlagSet("djob", flag.ContinueOnError)

    params := &CliParams{
        key:      new(string),
        cmdArgs:  []string{},
        logLevel: new(string),
        ttl:      new(int),
    }

    flags.StringVar(params.key, "k", "", "Lock key. Default: filename")
    flags.StringVar(params.logLevel, "l", "error", "Log level: debug|info|warn|error|fatal. Default: error")
    flags.IntVar(params.ttl, "t", 60, "Ttl in seconds. Default: 60")
    flags.Usage = func() {
        fmt.Fprint(os.Stderr, "Usage: djob [OPTIONS] -- [command with arguments and options] \n")
        flags.PrintDefaults()
    }
    if err := flags.Parse(args); err != nil {
        return nil, flags, fmt.Errorf("parse flags: %v", err)
    }

    for k, arg := range args {
        if arg == "--" {
            params.cmdArgs = args[k+1:]
            break
        }
    }

    if len(params.cmdArgs) == 0 {
        return nil, flags, errors.New("command is required")
    }

    if len(*params.key) == 0 {
        params.key = &params.cmdArgs[0]
    }

    return params, flags, nil
}

To work with arguments I need to create flagset. There may be a few of these sets, but I’m good with one now.
I’ve identified the flags I wrote about earlier. Define them the types and pass the address of the variable (the field in the structure) where the parsed value should be placed. In order to make it easier for the user to understand how to use the tool I defined the function Usage. It prints in the output information, as well as descriptions of all flags and their default values.
Then I need to call the flagset method Parse([]string). This method gets the user’s arguments on the input. Next it will be os.Args. I deliberately made this feature independent of os.Args for the convenience of testing. Flags that I have not specified in the configuration will be ignored.
Next I found the flag -- and put everything after it as there is in the cmdArgs parameter. This is the command that the user wants to run. In addition, I made some minimal validation and at the output I got the structure command to run and parameters.

Logging

Before moving to testing I will talk about logging. Logging is another component that matters for testing implementation. In any case, it is required in almost any program, including the utility I describe.
I used the library logrus, because it has the ability to intercept logs with hooks. And I used this feature to read logs during testing. At the same time, the code of the main function itself does not have to adapt to tests.

type testHook struct {
    channel chan *log.Entry
}

func (hook *testHook) Fire(entry *log.Entry) error {
    hook.channel <- entry
    return nil
}

func (hook *testHook) Levels() []log.Level {
    return log.AllLevels
}

func newTestLogger() (*log.Logger, chan *log.Entry) {
    logger := log.New()
    logger.Out = ioutil.Discard
    logger.Level = log.DebugLevel

    entryChan := make(chan *log.Entry, 0)

    logger.Hooks.Add(&testHook{channel: entryChan})

    return logger, entryChan
}

I created a new logger with the output turned off and added a testHook to it. It instantly sends the logger to the channel instead of putting it into the console. Not only the messages themselves, but the whole structure that contains the full information, including the levels of the log entry, time, etc. It is all very convenient.

Main function

Now I can go directly to the main function. In it, I need to call parseCliParams and pass the os.Args there, and then start the process as I described in the previous article.
However, I have two obstacles:

I can’t test main function
I don’t have the ability to pass various arguments in tests other than os.Args (or it will be very ugly and unreliable) So I implemented Exec():

func Exec(osArgs []string, logger *log.Logger) int {

    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    params, flags, err := parseCliParams(osArgs)
    if err != nil {
        logger.Errorf("parse arguments: %v \n", err.Error())
        flags.Usage()
        return 2
    }

    lvl, err := logrus.ParseLevel(*params.logLevel)
    if err != nil {
        logger.Fatalf("ivalid log level %s \n", *params.logLevel)
    }
    logger.SetLevel(lvl)

    memLock := internal.NewMemoryLock()

    locked, err := memLock.Lock(ctx, *params.key, *params.clientId)
    if err != nil {
        logger.Errorf("Didn't lock: %v \n", err)
    }
    if !locked {
        logger.Info("Already locked \n")
        return 2
    }

    cmd, err := internal.ExecCmd(ctx, params.cmdArgs)
    if err != nil {
        logger.Errorf("exec command: %v \n", err)
        return 1
    }
    logger.Info("started")

    defer func() {
        unlocked, err := memLock.UnLock(ctx, *params.key, *params.clientId)
        if err != nil {
            logger.Errorf("can't unlock: %v", err)
        }
        if !unlocked {
            logger.Info("Already unlocked")
        }
    }()

    exitCode, err := internal.WaitCmd(ctx, cmd)
    if err != nil {
        logger.Fatalf("wait command: %v", err)
        return 1
    }

    return exitCode
}

memLock := internal.NewMemoryLock() This is a simple implementation of locks in memory. In this example, it is only necessary for the code to work. More locks I will consider in the following articles and now their implementation does not matter.

cmd, err := internal.ExecCmd(ctx, params.cmdArgs) Wrap on the function of starting a new process, which I discussed in the last article. It is not necessary to know the content of this function to understand the material of this article.

exitCode, err := internal.WaitCmd(ctx, cmd) This function waits for the process running with ExecCmd to complete. It returns the response code. In this case too, it is not necessary to know the content of the function.

I will go ahead and say that all the received code of the utility will be pushed on GitHub. But this is after I go through all the parts of the implementation step by step.

func main() {
    logger := log.New()
    os.Exit(Exec(os.Args, logger))
}

Testing

Finally I came to the most important thing for which this article - testing. I did all the necessary preparation.
I needed a script to run in the tests. It has to be a program that can run for a while and I need to be able to set that time.

#!/usr/bin/env bash

trap "echo SIGINT; exit" SIGINT

for (( i=1; i<$1; i++ ))
do
  echo "step $i"
  sleep $2
done

echo "finish"

In the first argument I set the number of iterations, and in the second time the iteration. Everything is simple and this should be enough.

Next is the test itself.

func TestRunJob(t *testing.T) {
    logger, entryChan := newTestLogger()
    testCases := []struct {
        name       string
        args       []string
        logEntries []*struct {
            level   log.Level
            message string
        }
    }{
        {
            name: "first",
            args: []string{"djob", "-k", "key", "-c", "client1", "l", "info", "--", "../tests/script.sh", "2", "2"},
            logEntries: []*struct {
                level   log.Level
                message string
            }{
                {
                    level:   log.InfoLevel,
                    message: "start",
                },
                {
                    level:   log.InfoLevel,
                    message: "finish",
                },
            },
        },
    }

    for _, tt := range testCases {
        t.Run(tt.name, func(t *testing.T) {
            ctx, cancel := context.WithCancel(context.Background())
            defer cancel()

            go func() {
                entryIndex := 0
                for {
                    select {
                    case <-ctx.Done():
                        return
                    case entry := <-entryChan:
                        if entryIndex > len(tt.logEntries) {
                            t.Errorf("got more than expected (%d) log messages", len(tt.logEntries))
                        }
                        if tt.logEntries[entryIndex].level != entry.Level {
                            t.Errorf("Expected level: %s, actual: %s", tt.logEntries[entryIndex].level, entry.Level)
                        }
                        if tt.logEntries[entryIndex].message != entry.Message {
                            t.Errorf("Expected log message: %s, actual: %s", tt.logEntries[entryIndex].message, entry.Message)
                        }
                        entryIndex++
                    }
                }
            }()

            cmd.Exec(tt.args, logger)
        })
    }
}

Let me describe the whole code step by step.

testCases := []struct {
    name       string
    args       []string
    logEntries []*struct {
        level   log.Level
        message string
    }
}

in logEntries I specified expected entries in logs. I am only interested in message and level.

go func() {
                entryIndex := 0
                for {
                    select {
                    case <-ctx.Done():
                        return
                    case entry := <-entryChan:
                        if entryIndex > len(tt.logEntries) {
                            t.Errorf("got more than expected (%d) log messages", len(tt.logEntries))
                        }
                        if tt.logEntries[entryIndex].level != entry.Level {
                            t.Errorf("Expected level: %s, actual: %s", tt.logEntries[entryIndex].level, entry.Level)
                        }
                        if tt.logEntries[entryIndex].message != entry.Message {
                            t.Errorf("Expected log message: %s, actual: %s", tt.logEntries[entryIndex].message, entry.Message)
                        }
                        entryIndex++
                    }
                }
            }()

In the test I run a go routine. It listens to the channel with logs. And then I just make a check whether the entry is in the log, which I expect or not.

cmd. Exec(tt.args, logger) I run the main function with the necessary arguments and configured logger.

Conclusion

I described how it is possible to organize the testing of the cli utility with the help of logger. Next you only need to make all the necessary test cases. It was the second article in the Distributed Locks series. Next, I’ll move on to the implementation of the locks themselves and I’ll put the result on GitHub.

Blog

Testing cli tool with logging

Anton Gubarev

Flags

Logging

Main function

Testing

Conclusion

Join Our Newsletter. No Spam, Only the good stuff.

Related