github.com/mvdan/u-root-coreutils@v0.0.0-20230122170626-c2eef2898555/cmds/core/scp/scp_test.go (about)

     1  // Copyright 2018 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"os"
    12  	"path"
    13  	"testing"
    14  )
    15  
    16  func TestScpSource(t *testing.T) {
    17  	var w bytes.Buffer
    18  	var r bytes.Buffer
    19  
    20  	tf, err := os.CreateTemp("", "TestScpSource")
    21  	if err != nil {
    22  		t.Fatalf("creating temp file: %v", err)
    23  	}
    24  	defer os.Remove(tf.Name())
    25  	tf.Write([]byte("test-file-contents"))
    26  
    27  	r.Write([]byte{0})
    28  	err = scpSource(&w, &r, tf.Name())
    29  	if err != nil {
    30  		t.Fatalf("error: %v", err)
    31  	}
    32  	expected := []byte(fmt.Sprintf("C0600 18 %s\ntest-file-contents", path.Base(tf.Name())))
    33  	expected = append(expected, 0)
    34  	if string(expected) != w.String() {
    35  		t.Fatalf("Got: %v\nExpected: %v", w.String(), string(expected))
    36  	}
    37  }
    38  
    39  func TestScpSink(t *testing.T) {
    40  	var w bytes.Buffer
    41  	var r bytes.Buffer
    42  
    43  	tf, err := os.CreateTemp("", "TestScpSink")
    44  	if err != nil {
    45  		t.Fatalf("creating temp file: %v", err)
    46  	}
    47  	defer os.Remove(tf.Name())
    48  
    49  	r.Write([]byte("C0600 18 test\ntest-file-contents"))
    50  	// Post IO-copy success status
    51  	r.Write([]byte{0})
    52  
    53  	err = scpSink(&w, &r, tf.Name())
    54  	if err != nil {
    55  		t.Fatalf("error: %v", err)
    56  	}
    57  
    58  	// 1: Initial SUCCESS post to start scp
    59  	// 2: Success opening file tf.Name()
    60  	// 3: Success writing file
    61  	expected := []byte{0, 0, 0}
    62  	if string(expected) != w.String() {
    63  		t.Fatalf("Got: %v\nExpected: %v", w.Bytes(), expected)
    64  	}
    65  
    66  	m := make([]byte, 18)
    67  	n, err := tf.Read(m)
    68  	if err != nil {
    69  		t.Fatalf("IO error: %v", err)
    70  	}
    71  	if n != 18 {
    72  		t.Fatalf("Expected 18 bytes, got %v", n)
    73  	}
    74  
    75  	// Ensure EOF
    76  	_, err = tf.Read(m)
    77  	if err != io.EOF {
    78  		t.Fatalf("Expected EOF, got %v", err)
    79  	}
    80  
    81  	if string(m) != "test-file-contents" {
    82  		t.Fatalf("Expected 'test-file-contents', got '%v'", string(m))
    83  	}
    84  }