cocos/agent/algorithm/python/python_test.go

// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package python

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
	"github.com/ultravioletrs/cocos/agent/events/mocks"
	"google.golang.org/grpc/metadata"
)

const runtime = "python3"

func TestPythonRunTimeToContext(t *testing.T) {
	ctx := context.Background()
	newCtx := PythonRunTimeToContext(ctx, runtime)

	md, ok := metadata.FromOutgoingContext(newCtx)
	if !ok {
		t.Fatal("Expected metadata in context")
	}

	values := md.Get(PyRuntimeKey)
	if len(values) != 1 || values[0] != runtime {
		t.Errorf("Expected runtime %s, got %v", runtime, values)
	}
}

func TestPythonRunTimeFromContext(t *testing.T) {
	ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(PyRuntimeKey, runtime))

	got := PythonRunTimeFromContext(ctx)
	if got != runtime {
		t.Errorf("Expected runtime %s, got %s", runtime, got)
	}
}

func TestNewAlgorithm(t *testing.T) {
	logger := &slog.Logger{}
	eventsSvc := new(mocks.Service)
	requirementsFile := "requirements.txt"
	algoFile := "algorithm.py"
	args := []string{"--arg1", "value1"}

	algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args)

	p, ok := algo.(*python)
	if !ok {
		t.Fatal("Expected *python type")
	}

	if p.runtime != runtime {
		t.Errorf("Expected runtime %s, got %s", runtime, p.runtime)
	}
	if p.requirementsFile != requirementsFile {
		t.Errorf("Expected requirementsFile %s, got %s", requirementsFile, p.requirementsFile)
	}
	if p.algoFile != algoFile {
		t.Errorf("Expected algoFile %s, got %s", algoFile, p.algoFile)
	}
	if len(p.args) != len(args) {
		t.Errorf("Expected %d args, got %d", len(args), len(p.args))
	}
}

func TestRun(t *testing.T) {
	tmpDir, err := os.MkdirTemp("", "python-test")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)

	scriptContent := []byte("print('Hello, World!')")
	scriptPath := filepath.Join(tmpDir, "test_script.py")
	if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
		t.Fatal(err)
	}

	eventsSvc := new(mocks.Service)

	var stdout, stderr bytes.Buffer

	algo := &python{
		algoFile: scriptPath,
		stderr:   io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
		stdout:   io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
		runtime:  "python3",
	}

	err = algo.Run()
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	expectedOutput := "Hello, World!\n"
	if !strings.Contains(stdout.String(), expectedOutput) {
		t.Errorf("Expected output to contain %q, got %q", expectedOutput, stdout.String())
	}
}

func TestRunWithRequirements(t *testing.T) {
	tmpDir, err := os.MkdirTemp("", "python-test")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)

	scriptContent := []byte("import requests\nprint(requests.__version__)")
	scriptPath := filepath.Join(tmpDir, "test_script.py")
	if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
		t.Fatal(err)
	}

	requirementsContent := []byte("requests==2.26.0")
	requirementsPath := filepath.Join(tmpDir, "requirements.txt")
	if err := os.WriteFile(requirementsPath, requirementsContent, 0o644); err != nil {
		t.Fatal(err)
	}

	eventsSvc := new(mocks.Service)

	var stdout, stderr bytes.Buffer

	algo := &python{
		algoFile:         scriptPath,
		requirementsFile: requirementsPath,
		stderr:           io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
		stdout:           io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
		runtime:          "python3",
	}

	err = algo.Run()
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}

	if !strings.Contains(stdout.String(), "2.26.0") {
		t.Errorf("Expected output to contain requests version 2.26.0, got %q", stdout.String())
	}
}