cost-model/pkg/cloud/aws/athenaintegration_coverage_test.go

package aws

import (
	"strings"
	"testing"

	"github.com/opencost/opencost/pkg/cloud"
)

func TestAthenaWhereUsage(t *testing.T) {
	// Regression test for https://github.com/opencost/opencost/issues/4022 -
	// AWS Marketplace subscription/fee charges (line_item_line_item_type = 'Fee')
	// must be included, but only when billed through AWS Marketplace, so that
	// non-Marketplace 'Fee' rows (e.g. Reserved Instance upfront purchases,
	// already captured via 'DiscountedUsage' amortization) are not swept in and
	// double counted.
	expected := "(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount' OR (line_item_line_item_type = 'Fee' AND bill_billing_entity = 'AWS Marketplace'))"
	if AthenaWhereUsage != expected {
		t.Errorf("AthenaWhereUsage = %v, want %v", AthenaWhereUsage, expected)
	}

	if !strings.Contains(AthenaWhereUsage, "line_item_line_item_type = 'Fee' AND bill_billing_entity = 'AWS Marketplace'") {
		t.Errorf("AthenaWhereUsage should include AWS Marketplace 'Fee' rows scoped to bill_billing_entity, got: %v", AthenaWhereUsage)
	}

	// A bare, unscoped 'Fee' disjunct would also match non-Marketplace fees like RI
	// upfront purchases, so it must never appear on its own.
	if strings.Contains(AthenaWhereUsage, "line_item_line_item_type = 'Fee')") {
		t.Errorf("AthenaWhereUsage should not include an unscoped 'Fee' clause, got: %v", AthenaWhereUsage)
	}
}

func TestAthenaWhereUsageBase(t *testing.T) {
	// AthenaWhereUsageBase must reference only mandatory CUR columns (never
	// bill_billing_entity, which CUR 2.0 exports can disable) so it is always a safe
	// fallback when that column is absent.
	expected := "(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount')"
	if AthenaWhereUsageBase != expected {
		t.Errorf("AthenaWhereUsageBase = %v, want %v", AthenaWhereUsageBase, expected)
	}

	if strings.Contains(AthenaWhereUsageBase, "bill_billing_entity") || strings.Contains(AthenaWhereUsageBase, "Fee") {
		t.Errorf("AthenaWhereUsageBase must not reference bill_billing_entity or 'Fee', got: %v", AthenaWhereUsageBase)
	}
}

func TestAthenaIntegration_GetWhereUsage(t *testing.T) {
	ai := &AthenaIntegration{}

	// Regression test for https://github.com/opencost/opencost/pull/4028#discussion -
	// bill_billing_entity is a CUR 2.0-disableable column, like resource_tags,
	// line_item_usage_account_name, and bill_payer_account_name elsewhere in this file.
	// GetWhereUsage must check allColumns before referencing it, or CUR exports that
	// omit the column would fail every query with COLUMN_NOT_FOUND instead of merely
	// missing Marketplace fees.
	t.Run("bill_billing_entity present", func(t *testing.T) {
		allColumns := map[string]bool{
			AthenaBillingEntityColumn: true,
		}
		got := ai.GetWhereUsage(allColumns)
		if got != AthenaWhereUsage {
			t.Errorf("GetWhereUsage() = %v, want AthenaWhereUsage (%v)", got, AthenaWhereUsage)
		}
		if !strings.Contains(got, "bill_billing_entity = 'AWS Marketplace'") {
			t.Errorf("GetWhereUsage() should include the Marketplace fee clause when bill_billing_entity is present, got: %v", got)
		}
	})

	t.Run("bill_billing_entity absent", func(t *testing.T) {
		allColumns := map[string]bool{
			"line_item_line_item_type": true,
		}
		got := ai.GetWhereUsage(allColumns)
		if got != AthenaWhereUsageBase {
			t.Errorf("GetWhereUsage() = %v, want AthenaWhereUsageBase (%v)", got, AthenaWhereUsageBase)
		}
		if strings.Contains(got, "bill_billing_entity") {
			t.Errorf("GetWhereUsage() must not reference bill_billing_entity when it is absent from allColumns, got: %v", got)
		}
	})

	t.Run("empty allColumns", func(t *testing.T) {
		got := ai.GetWhereUsage(map[string]bool{})
		if got != AthenaWhereUsageBase {
			t.Errorf("GetWhereUsage() = %v, want AthenaWhereUsageBase (%v)", got, AthenaWhereUsageBase)
		}
	})
}

func TestAthenaIntegration_GetListCostColumn(t *testing.T) {
	ai := &AthenaIntegration{}
	expected := "SUM(CASE line_item_line_item_type WHEN 'EdpDiscount' THEN 0 WHEN 'PrivateRateDiscount' THEN 0 ELSE line_item_unblended_cost END) as list_cost"
	actual := ai.GetListCostColumn()
	if actual != expected {
		t.Errorf("GetListCostColumn() = %v, want %v", actual, expected)
	}
}

func TestAthenaIntegration_GetNetCostColumn(t *testing.T) {
	ai := &AthenaIntegration{}

	// Test case where net pricing column exists
	allColumnsWithNet := map[string]bool{
		"line_item_net_unblended_cost": true,
	}
	expectedWithNet := "SUM(COALESCE(line_item_net_unblended_cost, line_item_unblended_cost, 0)) as net_cost"
	actualWithNet := ai.GetNetCostColumn(allColumnsWithNet)
	if actualWithNet != expectedWithNet {
		t.Errorf("GetNetCostColumn() with net pricing = %v, want %v", actualWithNet, expectedWithNet)
	}

	// Test case where net pricing column doesn't exist
	allColumnsWithoutNet := map[string]bool{
		"line_item_unblended_cost": true,
	}
	expectedWithoutNet := "SUM(line_item_unblended_cost) as net_cost"
	actualWithoutNet := ai.GetNetCostColumn(allColumnsWithoutNet)
	if actualWithoutNet != expectedWithoutNet {
		t.Errorf("GetNetCostColumn() without net pricing = %v, want %v", actualWithoutNet, expectedWithoutNet)
	}
}

func TestAthenaIntegration_GetAmortizedCostColumn(t *testing.T) {
	ai := &AthenaIntegration{}
	allColumns := map[string]bool{
		"reservation_effective_cost":               true,
		"savings_plan_savings_plan_effective_cost": true,
		"line_item_unblended_cost":                 true,
	}

	result := ai.GetAmortizedCostColumn(allColumns)
	if !strings.Contains(result, "SUM(") || !strings.Contains(result, " as amortized_cost") {
		t.Errorf("GetAmortizedCostColumn() should return a SUM expression with amortized_cost alias, got: %v", result)
	}
}

func TestAthenaIntegration_GetAmortizedNetCostColumn(t *testing.T) {
	ai := &AthenaIntegration{}

	// Test case where net pricing columns exist
	allColumnsWithNet := map[string]bool{
		"line_item_net_unblended_cost":                 true,
		"reservation_net_effective_cost":               true,
		"savings_plan_net_savings_plan_effective_cost": true,
		"line_item_unblended_cost":                     true,
	}
	resultWithNet := ai.GetAmortizedNetCostColumn(allColumnsWithNet)
	if !strings.Contains(resultWithNet, "SUM(") || !strings.Contains(resultWithNet, " as amortized_net_cost") {
		t.Errorf("GetAmortizedNetCostColumn() with net pricing should return a SUM expression with amortized_net_cost alias, got: %v", resultWithNet)
	}

	// Test case where net pricing columns don't exist
	allColumnsWithoutNet := map[string]bool{
		"reservation_effective_cost":               true,
		"savings_plan_savings_plan_effective_cost": true,
		"line_item_unblended_cost":                 true,
	}
	resultWithoutNet := ai.GetAmortizedNetCostColumn(allColumnsWithoutNet)
	if !strings.Contains(resultWithoutNet, "SUM(") || !strings.Contains(resultWithoutNet, " as amortized_net_cost") {
		t.Errorf("GetAmortizedNetCostColumn() without net pricing should return a SUM expression with amortized_net_cost alias, got: %v", resultWithoutNet)
	}
}

func TestAthenaIntegration_GetAmortizedCostCase(t *testing.T) {
	ai := &AthenaIntegration{}

	// Test case where RI and SP pricing columns exist
	allColumnsWithRIAndSP := map[string]bool{
		"reservation_effective_cost":               true,
		"savings_plan_savings_plan_effective_cost": true,
		"line_item_unblended_cost":                 true,
	}
	resultWithRIAndSP := ai.GetAmortizedCostCase(allColumnsWithRIAndSP)
	if !strings.Contains(resultWithRIAndSP, "CASE line_item_line_item_type") ||
		!strings.Contains(resultWithRIAndSP, "DiscountedUsage") ||
		!strings.Contains(resultWithRIAndSP, "SavingsPlanCoveredUsage") {
		t.Errorf("GetAmortizedCostCase() with RI and SP should contain CASE statement with DiscountedUsage and SavingsPlanCoveredUsage, got: %v", resultWithRIAndSP)
	}

	// Test case where neither RI nor SP pricing columns exist
	allColumnsWithoutRIOrSP := map[string]bool{
		"line_item_unblended_cost": true,
	}
	resultWithoutRIOrSP := ai.GetAmortizedCostCase(allColumnsWithoutRIOrSP)
	expectedWithoutRIOrSP := "line_item_unblended_cost"
	if resultWithoutRIOrSP != expectedWithoutRIOrSP {
		t.Errorf("GetAmortizedCostCase() without RI or SP should return line_item_unblended_cost, got: %v, want: %v", resultWithoutRIOrSP, expectedWithoutRIOrSP)
	}
}

func TestAthenaIntegration_GetAmortizedNetCostCase(t *testing.T) {
	ai := &AthenaIntegration{}

	// Test case where net RI and SP pricing columns exist
	allColumnsWithNetRIAndSP := map[string]bool{
		"reservation_net_effective_cost":               true,
		"savings_plan_net_savings_plan_effective_cost": true,
		"line_item_net_unblended_cost":                 true,
		"line_item_unblended_cost":                     true,
	}
	resultWithNetRIAndSP := ai.GetAmortizedNetCostCase(allColumnsWithNetRIAndSP)
	if !strings.Contains(resultWithNetRIAndSP, "CASE line_item_line_item_type") ||
		!strings.Contains(resultWithNetRIAndSP, "DiscountedUsage") ||
		!strings.Contains(resultWithNetRIAndSP, "SavingsPlanCoveredUsage") {
		t.Errorf("GetAmortizedNetCostCase() with net RI and SP should contain CASE statement with DiscountedUsage and SavingsPlanCoveredUsage, got: %v", resultWithNetRIAndSP)
	}

	// Test case where neither net RI nor net SP pricing columns exist
	allColumnsWithoutNetRIOrSP := map[string]bool{
		"line_item_net_unblended_cost": true,
		"line_item_unblended_cost":     true,
	}
	resultWithoutNetRIOrSP := ai.GetAmortizedNetCostCase(allColumnsWithoutNetRIOrSP)
	expectedStr := "COALESCE(line_item_net_unblended_cost, line_item_unblended_cost, 0)"
	if resultWithoutNetRIOrSP != expectedStr {
		t.Errorf("GetAmortizedNetCostCase() without net RI or SP should return COALESCE expression, got: %v, want: %v", resultWithoutNetRIOrSP, expectedStr)
	}
}

func TestAthenaIntegration_RemoveColumnAliases(t *testing.T) {
	ai := &AthenaIntegration{}
	columns := []string{
		"column1 as alias1",
		"column2",
		"column3 as alias3",
		"column4",
	}

	ai.RemoveColumnAliases(columns)

	if columns[0] != "column1" {
		t.Errorf("RemoveColumnAliases() should remove alias from 'column1 as alias1', got: %v", columns[0])
	}
	if columns[1] != "column2" {
		t.Errorf("RemoveColumnAliases() should not modify 'column2', got: %v", columns[1])
	}
	if columns[2] != "column3" {
		t.Errorf("RemoveColumnAliases() should remove alias from 'column3 as alias3', got: %v", columns[2])
	}
	if columns[3] != "column4" {
		t.Errorf("RemoveColumnAliases() should not modify 'column4', got: %v", columns[3])
	}
}

func TestAthenaIntegration_ConvertLabelToAWSTag(t *testing.T) {
	ai := &AthenaIntegration{}

	// Test case where label already has prefix
	labelWithPrefix := "resource_tags_user_test_label"
	resultWithPrefix := ai.ConvertLabelToAWSTag(labelWithPrefix)
	if resultWithPrefix != labelWithPrefix {
		t.Errorf("ConvertLabelToAWSTag() should return label unchanged if it already has prefix, got: %v, want: %v", resultWithPrefix, labelWithPrefix)
	}

	// Test case where label needs prefix
	labelWithoutPrefix := "test.label/with:characters-here"
	resultWithoutPrefix := ai.ConvertLabelToAWSTag(labelWithoutPrefix)
	expectedWithoutPrefix := "resource_tags_user_test_label_with_characters_here"
	if resultWithoutPrefix != expectedWithoutPrefix {
		t.Errorf("ConvertLabelToAWSTag() should add prefix and replace characters, got: %v, want: %v", resultWithoutPrefix, expectedWithoutPrefix)
	}
}

func TestAthenaIntegration_GetIsKubernetesColumn(t *testing.T) {
	ai := &AthenaIntegration{}

	// Test with some tag columns present CUR 1.0
	allColumnsCur10 := map[string]bool{
		"resource_tags_user_eks_cluster_name":             true,
		"resource_tags_user_alpha_eksctl_io_cluster_name": true,
		"resource_tags_user_kubernetes_io_service_name":   true,
		"some_other_column":                               true,
	}

	result := ai.GetIsKubernetesColumn(allColumnsCur10)
	if !strings.Contains(result, "line_item_product_code = 'AmazonEKS'") {
		t.Errorf("GetIsKubernetesColumn() should always include EKS check, got: %v", result)
	}
	if !strings.Contains(result, "resource_tags_user_eks_cluster_name <> ''") {
		t.Errorf("GetIsKubernetesColumn() should include checks for tag columns, got: %v", result)
	}
	if !strings.Contains(result, " as is_kubernetes") {
		t.Errorf("GetIsKubernetesColumn() should alias result as is_kubernetes, got: %v", result)
	}

	// Test with some tag columns present CUR 2.0
	allColumnsCur20 := map[string]bool{
		"resource_tags":     true,
		"some_other_column": true,
	}

	result = ai.GetIsKubernetesColumn(allColumnsCur20)
	if !strings.Contains(result, "line_item_product_code = 'AmazonEKS'") {
		t.Errorf("GetIsKubernetesColumn() should always include EKS check, got: %v", result)
	}
	if !strings.Contains(result, "COALESCE(resource_tags['user_eks_cluster_name'], '') <> ''") {
		t.Errorf("GetIsKubernetesColumn() should include checks for tag columns, got: %v", result)
	}
	if !strings.Contains(result, " as is_kubernetes") {
		t.Errorf("GetIsKubernetesColumn() should alias result as is_kubernetes, got: %v", result)
	}
}

func TestAthenaQuerier_GetStatus(t *testing.T) {
	aq := &AthenaQuerier{}

	// Test initial status
	status := aq.GetStatus()
	if status.String() != cloud.InitialStatus.String() {
		t.Errorf("GetStatus() should return InitialStatus for uninitialized querier, got: %v", status)
	}

	// Test setting a specific status
	aq.ConnectionStatus = cloud.SuccessfulConnection
	status = aq.GetStatus()
	if status != cloud.SuccessfulConnection {
		t.Errorf("GetStatus() should return set status, got: %v", status)
	}
}

func TestAthenaQuerier_Equals(t *testing.T) {
	aq1 := &AthenaQuerier{
		AthenaConfiguration: AthenaConfiguration{
			Bucket:   "bucket1",
			Region:   "region1",
			Database: "database1",
			Table:    "table1",
			Account:  "account1",
			Authorizer: &AccessKey{
				ID:     "id1",
				Secret: "secret1",
			},
		},
	}

	aq2 := &AthenaQuerier{
		AthenaConfiguration: AthenaConfiguration{
			Bucket:   "bucket1",
			Region:   "region1",
			Database: "database1",
			Table:    "table1",
			Account:  "account1",
			Authorizer: &AccessKey{
				ID:     "id1",
				Secret: "secret1",
			},
		},
	}

	aq3 := &AthenaQuerier{
		AthenaConfiguration: AthenaConfiguration{
			Bucket:   "bucket2", // Different bucket
			Region:   "region1",
			Database: "database1",
			Table:    "table1",
			Account:  "account1",
			Authorizer: &AccessKey{
				ID:     "id1",
				Secret: "secret1",
			},
		},
	}

	// Test equality
	if !aq1.Equals(aq2) {
		t.Errorf("Equals() should return true for identical configurations")
	}

	// Test inequality
	if aq1.Equals(aq3) {
		t.Errorf("Equals() should return false for different configurations")
	}

	// Test comparison with non-AthenaQuerier
	accessKey := &AccessKey{
		ID:     "id1",
		Secret: "secret1",
	}
	if aq1.Equals(accessKey) {
		t.Errorf("Equals() should return false when comparing with different type")
	}
}