zzet--gortex/internal/parser/languages/python_test.go

package languages

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zzet/gortex/internal/graph"
)

func TestPyExtractor_Function(t *testing.T) {
	src := []byte(`def greet(name):
    return f"Hello {name}"
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	funcs := nodesOfKind(result.Nodes, graph.KindFunction)
	require.Len(t, funcs, 1)
	assert.Equal(t, "greet", funcs[0].Name)
}

func TestPyExtractor_Class(t *testing.T) {
	src := []byte(`class UserService:
    def __init__(self):
        self.users = []

    def get_user(self, user_id):
        return None
`)
	e := NewPythonExtractor()
	result, err := e.Extract("service.py", src)
	require.NoError(t, err)

	types := nodesOfKind(result.Nodes, graph.KindType)
	require.Len(t, types, 1)
	assert.Equal(t, "UserService", types[0].Name)

	// Class methods are extracted as KindMethod, not KindFunction.
	methods := nodesOfKind(result.Nodes, graph.KindMethod)
	assert.Len(t, methods, 2) // __init__, get_user

	// No top-level functions.
	funcs := nodesOfKind(result.Nodes, graph.KindFunction)
	assert.Len(t, funcs, 0)

	// EdgeMemberOf edges link methods to the class.
	memberEdges := edgesOfKind(result.Edges, graph.EdgeMemberOf)
	assert.Len(t, memberEdges, 2)
	for _, e := range memberEdges {
		assert.Equal(t, "service.py::UserService", e.To)
	}
}

func TestPyExtractor_Imports(t *testing.T) {
	src := []byte(`import os
from pathlib import Path
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	imports := edgesOfKind(result.Edges, graph.EdgeImports)
	require.Len(t, imports, 2)
}

func TestPyExtractor_TypeEnv_TypeHint(t *testing.T) {
	src := []byte(`
class UserService:
    def save(self):
        pass

def main():
    svc: UserService = get_service()
    svc.save()
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	calls := edgesOfKind(result.Edges, graph.EdgeCalls)
	var saveCall *graph.Edge
	for _, c := range calls {
		if strings.HasSuffix(c.To, "save") {
			saveCall = c
			break
		}
	}
	require.NotNil(t, saveCall, "expected a call edge to save")
	require.NotNil(t, saveCall.Meta, "expected Meta on save call edge")
	assert.Equal(t, "UserService", saveCall.Meta["receiver_type"])
}

func TestPyExtractor_TypeEnv_ClassConstructor(t *testing.T) {
	src := []byte(`
class Client:
    def connect(self):
        pass

def main():
    client = Client()
    client.connect()
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	calls := edgesOfKind(result.Edges, graph.EdgeCalls)
	var connectCall *graph.Edge
	for _, c := range calls {
		if strings.HasSuffix(c.To, "connect") {
			connectCall = c
			break
		}
	}
	require.NotNil(t, connectCall)
	require.NotNil(t, connectCall.Meta)
	assert.Equal(t, "Client", connectCall.Meta["receiver_type"])
}

func TestPyExtractor_TypeEnv_Chain(t *testing.T) {
	src := []byte(`
class Connection:
    def query(self) -> Result:
        return Result()

class Result:
    def first(self) -> User:
        return User()

class User:
    def save(self):
        pass

def main():
    conn = Connection()
    conn.query().first().save()
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	calls := edgesOfKind(result.Edges, graph.EdgeCalls)
	var saveCall *graph.Edge
	for _, c := range calls {
		if strings.HasSuffix(c.To, "save") {
			saveCall = c
			break
		}
	}
	require.NotNil(t, saveCall, "expected a call edge to save")
	require.NotNil(t, saveCall.Meta, "expected Meta on chained save call edge")
	assert.Equal(t, "User", saveCall.Meta["receiver_type"])
}

func TestPyExtractor_TypeEnv_Unknown(t *testing.T) {
	src := []byte(`
def get_service():
    return None

def main():
    svc = get_service()
    svc.process()
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	calls := edgesOfKind(result.Edges, graph.EdgeCalls)
	var processCall *graph.Edge
	for _, c := range calls {
		if strings.HasSuffix(c.To, "process") {
			processCall = c
			break
		}
	}
	require.NotNil(t, processCall)
	assert.NotContains(t, processCall.Meta, "receiver_type", "unknown type should not produce a receiver_type hint")
}

func TestPythonExtractor_FastAPIDepends(t *testing.T) {
	// Depends(target) in a parameter default (or Annotated[T, Depends(target)])
	// should produce a direct call edge from the handler to target, not
	// just to the generic Depends function. Without this pass,
	// callers(target) is empty for any DI-only factory.
	src := []byte(`
from fastapi import Depends
from typing import Annotated

def get_settings():
    return {"db": "x"}

def handler(settings: Annotated[dict, Depends(get_settings)]):
    return settings
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	var found *graph.Edge
	for _, c := range edgesOfKind(result.Edges, graph.EdgeCalls) {
		if c.Meta == nil {
			continue
		}
		if v, _ := c.Meta["via"].(string); v == "fastapi.Depends" {
			if strings.HasSuffix(c.To, "get_settings") {
				found = c
				break
			}
		}
	}
	require.NotNil(t, found, "expected a fastapi.Depends edge to get_settings")
	assert.Equal(t, "app.py::handler", found.From)
}

func TestPythonExtractor_DependsOnlyOnIdentifierArg(t *testing.T) {
	// Depends() with a non-identifier argument (lambda, attribute access)
	// shouldn't produce a bogus edge — we can only statically resolve
	// plain identifier targets.
	src := []byte(`
from fastapi import Depends

def handler(x = Depends(lambda: 42), y = Depends(obj.method)):
    return x
`)
	e := NewPythonExtractor()
	result, err := e.Extract("app.py", src)
	require.NoError(t, err)

	for _, c := range edgesOfKind(result.Edges, graph.EdgeCalls) {
		if c.Meta == nil {
			continue
		}
		if v, _ := c.Meta["via"].(string); v == "fastapi.Depends" {
			t.Fatalf("unexpected fastapi.Depends edge: %+v", c)
		}
	}
}

func TestPythonExtractor_DocAndVisibility(t *testing.T) {
	src := []byte(`def foo():
    """Foo does the thing.

    Some details.
    """
    return 1


def _helper():
    """Internal helper."""
    return 2


class Server:
    """The server."""

    def start(self):
        """Start it."""
        return None

    def _stop(self):
        return None
`)
	e := NewPythonExtractor()
	result, err := e.Extract("server.py", src)
	require.NoError(t, err)

	byID := map[string]*graph.Node{}
	for _, n := range result.Nodes {
		byID[n.ID] = n
	}

	foo := byID["server.py::foo"]
	require.NotNil(t, foo)
	if foo.Meta["doc"] != "Foo does the thing." {
		t.Fatalf("foo.doc = %q", foo.Meta["doc"])
	}
	if foo.Meta["visibility"] != "public" {
		t.Fatalf("foo.vis = %q", foo.Meta["visibility"])
	}

	helper := byID["server.py::_helper"]
	require.NotNil(t, helper)
	if helper.Meta["visibility"] != "private" {
		t.Fatalf("_helper.vis = %q", helper.Meta["visibility"])
	}

	server := byID["server.py::Server"]
	require.NotNil(t, server)
	if server.Meta["doc"] != "The server." {
		t.Fatalf("Server.doc = %q", server.Meta["doc"])
	}
	if server.Meta["visibility"] != "public" {
		t.Fatalf("Server.vis = %q", server.Meta["visibility"])
	}

	start := byID["server.py::Server.start"]
	require.NotNil(t, start)
	if start.Meta["doc"] != "Start it." {
		t.Fatalf("Server.start.doc = %q", start.Meta["doc"])
	}

	stop := byID["server.py::Server._stop"]
	require.NotNil(t, stop)
	if stop.Meta["visibility"] != "private" {
		t.Fatalf("Server._stop.vis = %q", stop.Meta["visibility"])
	}
}

func TestPythonExtractor_AnnotationEdges(t *testing.T) {
	src := []byte(`from app import app
from functools import lru_cache

@app.route("/users/:id")
@lru_cache
def get_user(id):
    return {}

@dataclass
class User:
    name: str

@deprecated
def old_helper():
    pass
`)
	e := NewPythonExtractor()
	result, err := e.Extract("svc.py", src)
	require.NoError(t, err)

	annNames := map[string]bool{}
	for _, n := range result.Nodes {
		if v, _ := n.Meta["kind"].(string); v == "annotation" {
			annNames[n.Name] = true
		}
	}
	for _, want := range []string{"app.route", "lru_cache", "dataclass", "deprecated"} {
		if !annNames[want] {
			t.Fatalf("missing annotation node %q (got %v)", want, annNames)
		}
	}

	edges := map[string][]string{}
	argsByEdge := map[string]string{}
	for _, e := range result.Edges {
		if e.Kind != graph.EdgeAnnotated {
			continue
		}
		edges[e.From] = append(edges[e.From], e.To)
		if v, ok := e.Meta["args"].(string); ok {
			argsByEdge[e.From+"->"+e.To] = v
		}
	}

	if !pyContains(edges["svc.py::get_user"], "annotation::python::app.route") {
		t.Fatalf("missing app.route edge on get_user, got %v", edges["svc.py::get_user"])
	}
	if argsByEdge["svc.py::get_user->annotation::python::app.route"] != `"/users/:id"` {
		t.Fatalf("app.route args = %q", argsByEdge["svc.py::get_user->annotation::python::app.route"])
	}
	if !pyContains(edges["svc.py::User"], "annotation::python::dataclass") {
		t.Fatalf("missing dataclass edge on User, got %v", edges["svc.py::User"])
	}
}

func pyContains(haystack []string, needle string) bool {
	for _, h := range haystack {
		if h == needle {
			return true
		}
	}
	return false
}

func TestPythonExtractor_ThrowsRaise(t *testing.T) {
	src := []byte(`def parse(s):
    if not s:
        raise ValueError("empty")
    if len(s) > 10:
        raise errors.TooLong()
    return int(s)


def silent():
    try:
        x = 1
    except Exception:
        pass
`)
	e := NewPythonExtractor()
	result, err := e.Extract("p.py", src)
	require.NoError(t, err)

	throws := edgesOfKind(result.Edges, graph.EdgeThrows)
	parseTargets := map[string]bool{}
	for _, e := range throws {
		if e.From == "p.py::parse" {
			parseTargets[e.To] = true
		}
	}
	assert.True(t, parseTargets["unresolved::ValueError"], "ValueError missing")
	assert.True(t, parseTargets["unresolved::TooLong"], "TooLong (attribute access) missing")

	for _, e := range throws {
		if e.From == "p.py::silent" {
			t.Fatalf("silent() doesn't raise; got %v", e)
		}
	}
}

func TestPythonExtractor_ImportNodes(t *testing.T) {
	// Note: relative imports (`from .local import x`) aren't covered
	// by the current import query — pre-existing limitation. The
	// import-node promotion only sees absolute paths today.
	src := []byte(`import os
import requests
from app import service
`)
	e := NewPythonExtractor()
	result, err := e.Extract("svc.py", src)
	require.NoError(t, err)

	imports := nodesOfKind(result.Nodes, graph.KindImport)
	require.GreaterOrEqual(t, len(imports), 3)

	byID := map[string]*graph.Node{}
	for _, n := range imports {
		byID[n.ID] = n
	}

	osNode := byID["svc.py::import::os"]
	require.NotNil(t, osNode)
	assert.Equal(t, false, osNode.Meta["is_external"], "os is stdlib")

	req := byID["svc.py::import::requests"]
	require.NotNil(t, req)
	assert.Equal(t, true, req.Meta["is_external"])
}

// Decorated methods (@property, @staticmethod, @contextmanager, ...) are
// wrapped in a decorated_definition node between the def and the class
// body. They must still classify as class members: two same-named
// decorated properties on two classes in one file used to collapse into
// bare colliding `<file>::<name>` IDs (httpx's Headers.encoding vs
// Response.encoding), so every consumer addressing the class-qualified
// ID silently landed on the wrong declaration.
func TestPyExtractor_DecoratedMethodsAreClassMembers(t *testing.T) {
	src := []byte(`class Headers:
    @property
    def encoding(self):
        return "ascii"

class Response:
    @property
    def encoding(self):
        return "utf-8"

    @staticmethod
    def build():
        return Response()

def caller():
    @wraps
    def inner():
        pass
    return inner
`)
	e := NewPythonExtractor()
	result, err := e.Extract("m.py", src)
	require.NoError(t, err)

	byID := map[string]*graph.Node{}
	for _, n := range result.Nodes {
		byID[n.ID] = n
	}

	he := byID["m.py::Headers.encoding"]
	require.NotNil(t, he, "decorated property must get a class-qualified ID (got IDs: %v)", pyNodeIDs(result.Nodes))
	assert.Equal(t, graph.KindMethod, he.Kind)
	assert.Equal(t, 3, he.StartLine)
	assert.Equal(t, "Headers", he.Meta["receiver"])

	re := byID["m.py::Response.encoding"]
	require.NotNil(t, re, "same-named decorated property on a second class must stay distinct")
	assert.Equal(t, graph.KindMethod, re.Kind)
	assert.Equal(t, 8, re.StartLine)
	assert.Equal(t, "Response", re.Meta["receiver"])

	rb := byID["m.py::Response.build"]
	require.NotNil(t, rb, "@staticmethod def must classify as a method")
	assert.Equal(t, graph.KindMethod, rb.Kind)

	// A decorated function nested inside a function is NOT a class
	// member — it stays in the free-function bucket.
	in := byID["m.py::inner"]
	require.NotNil(t, in)
	assert.Equal(t, graph.KindFunction, in.Kind)

	// member_of edges link each decorated method to its class.
	memberTo := map[string]string{}
	for _, ed := range edgesOfKind(result.Edges, graph.EdgeMemberOf) {
		memberTo[ed.From] = ed.To
	}
	assert.Equal(t, "m.py::Headers", memberTo["m.py::Headers.encoding"])
	assert.Equal(t, "m.py::Response", memberTo["m.py::Response.encoding"])
	assert.Equal(t, "m.py::Response", memberTo["m.py::Response.build"])
}

func pyNodeIDs(nodes []*graph.Node) []string {
	out := make([]string, 0, len(nodes))
	for _, n := range nodes {
		out = append(out, n.ID)
	}
	return out
}