// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Unit tests for the Reporter implementation.
package imgstatus

import (
	"errors"
	"policy_manager/dbus"
	"policy_manager/mock/mockdbus"
	"policy_manager/mock/mockpolicyenforcer"
	"policy_manager/mock/mocksysapi"
	"policy_manager/policymanagerproto"
	"policy_manager/sysapi"
	"reflect"
	"strings"
	"testing"

	"github.com/golang/mock/gomock"
	"github.com/golang/protobuf/proto"
)

const (
	// Mock LSB file.
	testLSB = `CHROMEOS_AUSERVER=https://tools.google.com/service/update2
                CHROMEOS_BOARD_APPID={76E245CF-C0D0-444D-BA50-36739C18EB00}
                CHROMEOS_DEVSERVER=
                CHROMEOS_RELEASE_APPID={90F229CE-83E2-4FAF-8479-E368A34938B1}
                CHROMEOS_RELEASE_BOARD=lakitu-signed-prempkeys
                CHROMEOS_RELEASE_BRANCH_NUMBER=0
                CHROMEOS_RELEASE_BUILD_NUMBER=16511
                CHROMEOS_RELEASE_BUILD_TYPE=Official Build
                CHROMEOS_RELEASE_CHROME_MILESTONE=89
                CHROMEOS_RELEASE_DESCRIPTION=16511.0.0 (Official Build) dev-channel lakitu
                CHROMEOS_RELEASE_NAME=Chrome OS
                CHROMEOS_RELEASE_PATCH_NUMBER=0
                CHROMEOS_RELEASE_TRACK=dev-channel
                CHROMEOS_RELEASE_VERSION=16511.0.0
                GOOGLE_RELEASE=16511.0.0
                HWID_OVERRIDE=LAKITU`

	testDevBuildLSB = `
                CHROMEOS_RELEASE_BOARD=lakitu
                CHROMEOS_DEVSERVER=http://foo.google.com:8080
                GOOGLE_RELEASE=7520.3.2015_10_12_1443
                CHROMEOS_RELEASE_BUILD_NUMBER=7520
                CHROMEOS_RELEASE_BRANCH_NUMBER=3
                CHROMEOS_RELEASE_CHROME_MILESTONE=89
                CHROMEOS_RELEASE_PATCH_NUMBER=2015_10_12_1443
                CHROMEOS_RELEASE_TRACK=developer-build
                CHROMEOS_RELEASE_DESCRIPTION=7520.3.2015_10_12_1443 (Developer Build)
                CHROMEOS_RELEASE_NAME=Chromium OS
                CHROMEOS_RELEASE_BUILD_TYPE=Developer Build
                CHROMEOS_RELEASE_VERSION=7520.3.2015_10_12_1443
                CHROMEOS_AUSERVER=http://foo.google.com:8080/update`

	// Dummy instance ID.
	testInstanceID = uint64(12031993)
)

// TestParseKeyValuePairs tests that we can correctly parse /etc/lsb-release
// into a map.
func TestParseKeyValuePairs(t *testing.T) {
	tests := []struct {
		name  string
		want  map[string]string
		input string
	}{
		{
			"ParseLSBFile",
			map[string]string{
				"CHROMEOS_AUSERVER":                 "https://tools.google.com/service/update2",
				"CHROMEOS_BOARD_APPID":              "{76E245CF-C0D0-444D-BA50-36739C18EB00}",
				"CHROMEOS_DEVSERVER":                "",
				"CHROMEOS_RELEASE_APPID":            "{90F229CE-83E2-4FAF-8479-E368A34938B1}",
				"CHROMEOS_RELEASE_BOARD":            "lakitu-signed-prempkeys",
				"CHROMEOS_RELEASE_BRANCH_NUMBER":    "0",
				"CHROMEOS_RELEASE_BUILD_NUMBER":     "16511",
				"CHROMEOS_RELEASE_BUILD_TYPE":       "Official Build",
				"CHROMEOS_RELEASE_CHROME_MILESTONE": "89",
				"CHROMEOS_RELEASE_DESCRIPTION":      "16511.0.0 (Official Build) dev-channel lakitu",
				"CHROMEOS_RELEASE_NAME":             "Chrome OS",
				"CHROMEOS_RELEASE_PATCH_NUMBER":     "0",
				"CHROMEOS_RELEASE_TRACK":            "dev-channel",
				"CHROMEOS_RELEASE_VERSION":          "16511.0.0",
				"GOOGLE_RELEASE":                    "16511.0.0",
				"HWID_OVERRIDE":                     "LAKITU",
			},
			testLSB,
		},
		{
			"PoorlyFormattedPairs",
			map[string]string{
				"apple_A":    "a",
				"BAnaNA_b_":  "b",
				"cantaloupe": "",
			},
			`
                                apple_A= a
                                BAnaNA_b_ =b
                                    cantaloupe=
                        `,
		},
	}

	for _, test := range tests {
		mockReader := strings.NewReader(test.input)
		pairs, err := parseKeyValuePairs(mockReader)
		if err != nil {
			t.Errorf("Test %s got unexpected error %v",
				test.name,
				err)
		} else if !reflect.DeepEqual(pairs, test.want) {
			t.Errorf("Test %s got %v, want %v",
				test.name,
				pairs,
				test.want)
		}
	}
}

// TestConvertOSReleaseChannel tests that we can correctly convert the string
// representation of release channel into the protobuf definition.
func TestConvertOSReleaseChannel(t *testing.T) {
	tests := []struct {
		name  string
		want  policymanagerproto.ReleaseChannel
		input string
	}{
		{
			"DevChannelTest",
			policymanagerproto.ReleaseChannel_DEV,
			"dev-channel",
		},
		{
			"BetaChannelTest",
			policymanagerproto.ReleaseChannel_BETA,
			"beta-channel",
		},
		{
			"StableChannelTest",
			policymanagerproto.ReleaseChannel_STABLE,
			"stable-channel",
		},
		{
			"DeveloperBuildTest",
			policymanagerproto.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED,
			"developer-build",
		},
		{
			"BadChannelTest",
			policymanagerproto.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED,
			"abc-chan",
		},
	}

	for _, test := range tests {
		result := convertOSReleaseChannel(test.input)
		if !reflect.DeepEqual(test.want, result) {
			t.Errorf("Test %s got %s, want %s",
				test.name,
				result.String(),
				test.want.String())
		}
	}
}

// TestParseOSInfoFromLSB tests that we can correctly get the OS version and
// release channel from a lsb-release file.
func TestParseOSInfoFromLSB(t *testing.T) {
	devChannel := policymanagerproto.ReleaseChannel_DEV
	unknownChannel := policymanagerproto.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED

	tests := []struct {
		name       string
		want       *policymanagerproto.OSVersion
		lsbContent string
		expectErr  bool
	}{
		{
			"GoodLSBFile",
			&policymanagerproto.OSVersion{
				VersionString: proto.String("16511.0.0"),
				Milestone:     proto.Uint32(uint32(89)),
				Channel:       &devChannel,
			},
			testLSB,
			false,
		},
		{
			"OnlyReleaseVersion",
			&policymanagerproto.OSVersion{
				VersionString: proto.String("16511.0.0"),
			},
			"CHROMEOS_RELEASE_VERSION=16511.0.0",
			false,
		},
		{
			"BadLSBFile",
			&policymanagerproto.OSVersion{},
			"",
			false,
		},
		{
			"DevBuildLSBFile",
			&policymanagerproto.OSVersion{
				VersionString: proto.String("7520.3.2015_10_12_1443"),
				Milestone:     proto.Uint32(uint32(89)),
				Channel:       &unknownChannel,
			},
			testDevBuildLSB,
			false,
		},
	}

	// Set up gomock controller.
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	for _, test := range tests {
		mockAPI := mocksysapi.NewMockAPIHandler(mockCtrl)
		mockAPI.EXPECT().ReadFile("/etc/lsb-release").Return([]byte(test.lsbContent), nil)

		result, err := parseOSInfoFromLSB(mockAPI)

		if err == nil && test.expectErr {
			t.Errorf("Test %s passed, want error", test.name)
		} else if err == nil && !proto.Equal(test.want, result) {
			t.Errorf("Test %s got %s, want %s",
				test.name,
				result.String(),
				test.want.String())
		} else if err != nil && !test.expectErr {
			t.Errorf("Test %s got unexpected error %v",
				test.name,
				err)
		}
	}
}

// TestGetStatusFromUpdateEngine tests that we can correctly get the update
// status from update-engine via dbus.
func TestGetStatusFromUpdateEngine(t *testing.T) {
	rebootStatus := policymanagerproto.Operation_UPDATED_NEED_REBOOT
	tests := []struct {
		name              string
		getStatusResponse *policymanagerproto.StatusResult
		expectedRet       *policymanagerproto.InstanceStatus
		expectedError     error
	}{
		{
			"GoodUpdateEngineOutput",
			&policymanagerproto.StatusResult{
				NewVersion:      proto.String("16550.0.0"),
				UpdateStatus:    &rebootStatus,
				LastCheckedTime: proto.Int64(1433971203),
			},
			&policymanagerproto.InstanceStatus{
				NewOsVersion: &policymanagerproto.OSVersion{
					VersionString: proto.String("16550.0.0"),
				},
				UpdateStatus:         &rebootStatus,
				UpdateCheckTimestamp: proto.Int64(1433971203),
			},
			nil,
		},
		{
			"BadUpdateEngineOutput",
			&policymanagerproto.StatusResult{
				NewVersion:      proto.String(""),
				UpdateStatus:    nil,
				LastCheckedTime: proto.Int64(0),
			},
			nil,
			errors.New("GetStatus error"),
		},
	}

	for _, test := range tests {
		// Set up gomock controller.
		mockCtrl := gomock.NewController(t)

		mockClient := mockdbus.NewMockUpdateEngineClient(mockCtrl)
		mockClient.EXPECT().GetStatus().Return(
			test.getStatusResponse,
			test.expectedError)
		result, err := getStatusFromUpdateEngine(mockClient)

		if err == nil && test.expectedError != nil {
			t.Errorf("Test %s passed, want error %v", test.name, test.expectedError)
		} else if err == nil && !proto.Equal(test.expectedRet, result) {
			t.Errorf("Test %s got %s, want %s",
				test.name,
				result.String(),
				test.expectedRet.String())
		} else if err != nil && test.expectedError == nil {
			t.Errorf("Test %s got unexpected error %v", test.name, err)
		}

		mockCtrl.Finish()
	}
}

// TestGetStatus tests if the Reporter is returning the current status protobuf
// using information provided by a mock sysapi.APIHandler and a
// mock dbus.UpdateEngineClient.
func TestGetStatus(t *testing.T) {
	devChannel := policymanagerproto.ReleaseChannel_DEV

	// Set up gomock controller.
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mockAPI := mocksysapi.NewMockAPIHandler(mockCtrl)
	mockAPI.EXPECT().ReadFile("/etc/lsb-release").Return([]byte(testLSB), nil)

	mockClient := mockdbus.NewMockUpdateEngineClient(mockCtrl)
	rebootStatus := policymanagerproto.Operation_UPDATED_NEED_REBOOT

	mockClient.EXPECT().GetStatus().Return(
		&policymanagerproto.StatusResult{
			NewVersion:      proto.String("16550.0.0"),
			UpdateStatus:    &rebootStatus,
			LastCheckedTime: proto.Int64(1433971203),
		}, nil)

	mockPolicyEnforcer := mockpolicyenforcer.NewMockPolicyEnforcer(mockCtrl)
	mockPolicyEnforcer.EXPECT().GetHealthMonitorStatus().Return(
		&policymanagerproto.HealthMonitorStatus{
			Logging:    proto.Bool(true),
			Monitoring: proto.Bool(false),
		}, nil)

	reporter := NewReporter(sysapi.APIHandler(mockAPI), dbus.UpdateEngineClient(mockClient), mockPolicyEnforcer, testInstanceID)

	// Construct the response we want.
	currentVersion := policymanagerproto.OSVersion{
		VersionString: proto.String("16511.0.0"),
		Channel:       &devChannel,
		Milestone:     proto.Uint32(uint32(89)),
	}

	newVersion := policymanagerproto.OSVersion{
		VersionString: proto.String("16550.0.0"),
	}

	want := &policymanagerproto.InstanceStatus{
		InstanceId:           proto.Uint64(testInstanceID),
		OsVersion:            &currentVersion,
		NewOsVersion:         &newVersion,
		UpdateStatus:         &rebootStatus,
		UpdateCheckTimestamp: proto.Int64(1433971203),
		HealthMonitorStatus: &policymanagerproto.HealthMonitorStatus{
			Logging:    proto.Bool(true),
			Monitoring: proto.Bool(false),
		},
	}

	if status, err := reporter.GetStatus(); err != nil {
		t.Fatal(err)
	} else if !proto.Equal(status, want) {
		t.Errorf("got: %s, want: %s",
			proto.MarshalTextString(status),
			proto.MarshalTextString(want))
	}
}

// TestGetStatusError tests if the Reporter is returning a status protobuf with
// error field set in case an error occured while getting instance status.
// It also checks that the status is never nil and contains a valid InstanceId
// field even in the event of an error.
func TestGetStatusError(t *testing.T) {
	tests := []struct {
		name             string
		lsbErr           error
		updateEngineErr  error
		healthMonitorErr error
	}{
		{
			"LSBError",
			errors.New("test lsb error"),
			nil,
			nil,
		},
		{
			"UpdateEngineError",
			nil,
			errors.New("update-engine error"),
			nil,
		},
		{
			"HealthMonitorClientError",
			nil,
			nil,
			errors.New("health monitor client error"),
		},
	}

	for _, test := range tests {
		// Set up gomock controller.
		mockCtrl := gomock.NewController(t)

		mockAPI := mocksysapi.NewMockAPIHandler(mockCtrl)
		mockAPI.EXPECT().ReadFile("/etc/lsb-release").
			Return([]byte{}, test.lsbErr).AnyTimes()

		mockClient := mockdbus.NewMockUpdateEngineClient(mockCtrl)
		errorStatus := policymanagerproto.Operation_ERROR

		// Will only call GetStatus() if we didn't get error from ReadFile().
		if test.lsbErr == nil {
			mockClient.EXPECT().GetStatus().Return(
				&policymanagerproto.StatusResult{
					NewVersion:      proto.String(""),
					UpdateStatus:    &errorStatus,
					LastCheckedTime: proto.Int64(0),
				},
				test.updateEngineErr)
		}

		mockPolicyEnforcer := mockpolicyenforcer.NewMockPolicyEnforcer(mockCtrl)
		if test.lsbErr == nil && test.updateEngineErr == nil {
			mockPolicyEnforcer.EXPECT().GetHealthMonitorStatus().Return(
				&policymanagerproto.HealthMonitorStatus{
					Logging:    proto.Bool(true),
					Monitoring: proto.Bool(false),
				}, test.healthMonitorErr)
		}

		reporter := NewReporter(sysapi.APIHandler(mockAPI), dbus.UpdateEngineClient(mockClient), mockPolicyEnforcer, testInstanceID)

		if status, err := reporter.GetStatus(); err == nil {
			t.Errorf("test %s passed, expected error", test.name)
		} else if status == nil {
			t.Errorf("test %s returned nil error", test.name)
		} else if status.GetInstanceId() != testInstanceID {
			t.Errorf("test %s returned status with invalid instance ID field",
				test.name)
		} else if status.Error == nil {
			t.Errorf("test %s returned status without error field set",
				test.name)
		} else if status.GetError() != err.Error() {
			t.Errorf("test %s returned status with error: %s, expected: %s",
				test.name, status.GetError(), err.Error())
		}

		mockCtrl.Finish()
	}
}