Skip to content

refactor: issue 16 #72

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wiebe-vandendriessche merged 1 commit into from
Oct 8, 2026
Merged

refactor: issue 16 #72

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 37 additions & 15 deletions pkg/aibomgen/generator/generator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,22 @@ type bomBuilder interface {
BuildDataset(builder.DatasetBuildContext) (*cdx.Component, error)
}

var newBOMBuilder = func() bomBuilder {
return builder.BOMBuilder{}
// generator holds the dependency factories; tests construct one with doubles
// instead of mutating package state.
type generator struct {
newBOMBuilder func() bomBuilder
newFetcherSet func(httpClient *http.Client, baseURL string) fetcherSet
newDummyFetcherSet func() fetcherSet
}

func defaultGenerator() *generator {
return &generator{
newBOMBuilder: func() bomBuilder { return builder.BOMBuilder{} },
newFetcherSet: liveFetcherSet,
newDummyFetcherSet: dummyFetcherSet,
}
}

// Fetcher factory functions for testing.
type fetcherSet struct {
modelAPI interface {
FetchRevision(id, revision string) (*fetcher.ModelAPIResponse, error)
Expand Down Expand Up @@ -78,7 +89,7 @@ func ParseModelRef(s string) (ModelRef, error) {
return ModelRef{ID: id, Revision: rev}, nil
}

var newFetcherSet = func(httpClient *http.Client, baseURL string) fetcherSet {
func liveFetcherSet(httpClient *http.Client, baseURL string) fetcherSet {
return fetcherSet{
modelAPI: &fetcher.ModelAPIFetcher{Client: httpClient, BaseURL: baseURL},
modelReadme: &fetcher.ModelReadmeFetcher{Client: httpClient, BaseURL: baseURL},
Expand All @@ -92,8 +103,7 @@ func newHTTPClient(opts GenerateOptions) *http.Client {
return fetcher.NewHFClient(opts.Timeout, opts.HFToken)
}

// Dummy fetcher factory for BuildDummyBOM testing.
var newDummyFetcherSet = func() fetcherSet {
func dummyFetcherSet() fetcherSet {
return fetcherSet{
modelAPI: &fetcher.DummyModelAPIFetcher{},
modelReadme: &fetcher.DummyModelReadmeFetcher{},
Expand Down Expand Up @@ -148,7 +158,11 @@ type GenerateOptions struct {
// BuildDummyBOM builds a single comprehensive dummy BOM with all fields populated.
// This is used in dummy mode for testing/demo purposes without scanning or fetching real data.
func BuildDummyBOM() ([]DiscoveredBOM, error) {
fetchers := newDummyFetcherSet()
return defaultGenerator().buildDummyBOM()
}

func (g *generator) buildDummyBOM() ([]DiscoveredBOM, error) {
fetchers := g.newDummyFetcherSet()

// Create a dummy discovery.
dummyDiscovery := scanner.Discovery{
Expand Down Expand Up @@ -184,15 +198,15 @@ func BuildDummyBOM() ([]DiscoveredBOM, error) {
SecurityTree: securityTree,
}

bomBuilder := newBOMBuilder()
bomBuilder := g.newBOMBuilder()
bom, err := bomBuilder.Build(bctx)
if err != nil {
return nil, err
}

// Build dataset components for any datasets referenced in the model's training metadata.
noProgress := func(ProgressEvent) {}
_, resolved := buildDatasetComponents(fetchers, bom, extractDatasetsFromModel(apiResp, readme), "dummy-org/dummy-model", noProgress)
_, resolved := buildDatasetComponents(fetchers, bomBuilder, bom, extractDatasetsFromModel(apiResp, readme), "dummy-org/dummy-model", noProgress)
finalizeModelBOM(bom, resolved)

return []DiscoveredBOM{
Expand All @@ -208,6 +222,10 @@ func BuildDummyBOM() ([]DiscoveredBOM, error) {
// When building a model, if datasets are referenced in the model's training metadata, builds dataset components too.
// Use opts.OnProgress to receive progress events; pass a nil callback to disable.
func BuildPerDiscovery(discoveries []scanner.Discovery, opts GenerateOptions) ([]DiscoveredBOM, error) {
return defaultGenerator().buildPerDiscovery(discoveries, opts)
}

func (g *generator) buildPerDiscovery(discoveries []scanner.Discovery, opts GenerateOptions) ([]DiscoveredBOM, error) {
if opts.Timeout <= 0 {
opts.Timeout = 10 * time.Second
}
Expand All @@ -219,8 +237,8 @@ func BuildPerDiscovery(discoveries []scanner.Discovery, opts GenerateOptions) ([

results := make([]DiscoveredBOM, 0, len(discoveries))

fetchers := newFetcherSet(newHTTPClient(opts), opts.BaseURL)
bomBuilder := newBOMBuilder()
fetchers := g.newFetcherSet(newHTTPClient(opts), opts.BaseURL)
bomBuilder := g.newBOMBuilder()

for i, d := range discoveries {
modelID := strings.TrimSpace(d.ID)
Expand Down Expand Up @@ -291,7 +309,7 @@ func BuildPerDiscovery(discoveries []scanner.Discovery, opts GenerateOptions) ([

progress(ProgressEvent{Type: EventBuildComplete, ModelID: label})

datasetCount, resolved := buildDatasetComponents(fetchers, bom, extractDatasetsFromModel(resp, readme), label, progress)
datasetCount, resolved := buildDatasetComponents(fetchers, bomBuilder, bom, extractDatasetsFromModel(resp, readme), label, progress)
finalizeModelBOM(bom, resolved)

progress(ProgressEvent{Type: EventModelComplete, ModelID: label, Datasets: datasetCount})
Expand Down Expand Up @@ -394,7 +412,7 @@ func extractDatasetsFromModel(modelResp *fetcher.ModelAPIResponse, readme *fetch
// redirects renamed datasets, several card names can resolve to one component; it is
// added once. Dataset references that fail to fetch (e.g. not on HuggingFace) are
// skipped here and kept as inline entries by builder.LinkDatasetRefs.
func buildDatasetComponents(fetchers fetcherSet, bom *cdx.BOM, datasets []string, modelID string, progress ProgressCallback) (int, map[string]string) {
func buildDatasetComponents(fetchers fetcherSet, bomBuilder bomBuilder, bom *cdx.BOM, datasets []string, modelID string, progress ProgressCallback) (int, map[string]string) {
count := 0
resolved := make(map[string]string)
for _, dsID := range datasets {
Expand All @@ -420,7 +438,7 @@ func buildDatasetComponents(fetchers fetcherSet, bom *cdx.BOM, datasets []string
Readme: dsReadme,
}

dsComp, err := newBOMBuilder().BuildDataset(dsCtx)
dsComp, err := bomBuilder.BuildDataset(dsCtx)
if err != nil {
continue
}
Expand All @@ -445,6 +463,10 @@ func buildDatasetComponents(fetchers fetcherSet, bom *cdx.BOM, datasets []string
// Each ID may carry a revision as "org/name@revision" (see ParseModelRef).
// Use opts.OnProgress to receive progress events; pass a nil callback to disable.
func BuildFromModelIDs(modelIDs []string, opts GenerateOptions) ([]DiscoveredBOM, error) {
return defaultGenerator().buildFromModelIDs(modelIDs, opts)
}

func (g *generator) buildFromModelIDs(modelIDs []string, opts GenerateOptions) ([]DiscoveredBOM, error) {
discoveries := make([]scanner.Discovery, 0, len(modelIDs))
for _, rawID := range modelIDs {
if strings.TrimSpace(rawID) == "" {
Expand All @@ -465,5 +487,5 @@ func BuildFromModelIDs(modelIDs []string, opts GenerateOptions) ([]DiscoveredBOM
Revision: ref.Revision,
})
}
return BuildPerDiscovery(discoveries, opts)
return g.buildPerDiscovery(discoveries, opts)
}
48 changes: 28 additions & 20 deletions pkg/aibomgen/generator/generator_refs_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,24 +33,24 @@ func hfLikeFetchers(api *mockModelAPIFetcher) fetcherSet {
}
}

func withFetchers(t *testing.T, fs fetcherSet) {
t.Helper()
orig := newFetcherSet
newFetcherSet = func(*http.Client, string) fetcherSet { return fs }
t.Cleanup(func() { newFetcherSet = orig })
func withFetchers(fs fetcherSet) *generator {
g := defaultGenerator()
g.newFetcherSet = func(*http.Client, string) fetcherSet { return fs }
return g
}

func TestBuildFromModelIDs_DatasetRefsResolve(t *testing.T) {
t.Parallel()
api := &mockModelAPIFetcher{fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) {
return &fetcher.ModelAPIResponse{
ID: "org/model",
SHA: "abc",
CardData: map[string]any{"datasets": []any{"wikipedia", "legacy-datasets/wikipedia", "missing"}},
}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

boms, err := BuildFromModelIDs([]string{"org/model"}, GenerateOptions{})
boms, err := g.buildFromModelIDs([]string{"org/model"}, GenerateOptions{})
if err != nil || len(boms) != 1 {
t.Fatalf("BuildFromModelIDs = %d boms, err %v", len(boms), err)
}
Expand All @@ -74,6 +74,7 @@ func TestBuildFromModelIDs_DatasetRefsResolve(t *testing.T) {
}

func TestBuildDummyBOM_RefsResolve(t *testing.T) {
t.Parallel()
boms, err := BuildDummyBOM()
if err != nil || len(boms) != 1 {
t.Fatalf("BuildDummyBOM = %d boms, err %v", len(boms), err)
Expand All @@ -84,12 +85,13 @@ func TestBuildDummyBOM_RefsResolve(t *testing.T) {
}

func TestBuildPerDiscovery_RefsResolveAndRevision(t *testing.T) {
t.Parallel()
api := &mockModelAPIFetcher{fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) {
return &fetcher.ModelAPIResponse{ID: "org/model", SHA: "abc", CardData: map[string]any{"datasets": "wikipedia"}}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

boms, err := BuildPerDiscovery([]scanner.Discovery{{ID: "org/model", Name: "org/model", Type: "model", Revision: "v2"}}, GenerateOptions{})
boms, err := g.buildPerDiscovery([]scanner.Discovery{{ID: "org/model", Name: "org/model", Type: "model", Revision: "v2"}}, GenerateOptions{})
if err != nil || len(boms) != 1 {
t.Fatalf("BuildPerDiscovery = %d boms, err %v", len(boms), err)
}
Expand All @@ -105,12 +107,13 @@ func TestBuildPerDiscovery_RefsResolveAndRevision(t *testing.T) {
}

func TestBuildFromModelIDs_Revisions(t *testing.T) {
t.Parallel()
api := &mockModelAPIFetcher{fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) {
return &fetcher.ModelAPIResponse{ID: id, SHA: "ABC"}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

boms, err := BuildFromModelIDs([]string{"org/model@v1.0", "org/model", "org/model@"}, GenerateOptions{})
boms, err := g.buildFromModelIDs([]string{"org/model@v1.0", "org/model", "org/model@"}, GenerateOptions{})
if err != nil {
t.Fatal(err)
}
Expand Down Expand Up @@ -138,6 +141,7 @@ func TestBuildFromModelIDs_Revisions(t *testing.T) {
}

func TestParseModelRef(t *testing.T) {
t.Parallel()
tests := []struct {
in string
want ModelRef
Expand Down Expand Up @@ -166,6 +170,7 @@ func TestParseModelRef(t *testing.T) {
}

func TestBuildFromModelIDs_InputsOutputs(t *testing.T) {
t.Parallel()
tags := map[string]string{
"google-bert/bert-base-uncased": "fill-mask",
"sentence-transformers/all-MiniLM-L6-v2": "sentence-similarity",
Expand All @@ -174,10 +179,10 @@ func TestBuildFromModelIDs_InputsOutputs(t *testing.T) {
api := &mockModelAPIFetcher{fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) {
return &fetcher.ModelAPIResponse{ID: id, SHA: "abc", PipelineTag: tags[id]}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

ids := []string{"google-bert/bert-base-uncased", "sentence-transformers/all-MiniLM-L6-v2", "openai/whisper-large-v3"}
boms, err := BuildFromModelIDs(ids, GenerateOptions{})
boms, err := g.buildFromModelIDs(ids, GenerateOptions{})
if err != nil || len(boms) != len(ids) {
t.Fatalf("BuildFromModelIDs = %d boms, err %v", len(boms), err)
}
Expand Down Expand Up @@ -213,6 +218,7 @@ func TestBuildFromModelIDs_InputsOutputs(t *testing.T) {
}

func TestBuildFromModelIDs_Pedigree(t *testing.T) {
t.Parallel()
lineage := map[string]*fetcher.ModelBaseModels{
"Qwen/Qwen2.5-7B-Instruct": {Relation: "finetune"},
"nvidia/Eagle2.5-8B": {Relation: "merge"},
Expand All @@ -228,9 +234,9 @@ func TestBuildFromModelIDs_Pedigree(t *testing.T) {
BaseModels: lineage[id],
}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

boms, err := BuildFromModelIDs([]string{"Qwen/Qwen2.5-7B-Instruct", "nvidia/Eagle2.5-8B"}, GenerateOptions{})
boms, err := g.buildFromModelIDs([]string{"Qwen/Qwen2.5-7B-Instruct", "nvidia/Eagle2.5-8B"}, GenerateOptions{})
if err != nil || len(boms) != 2 {
t.Fatalf("BuildFromModelIDs = %d boms, err %v", len(boms), err)
}
Expand Down Expand Up @@ -268,17 +274,18 @@ func TestBuildFromModelIDs_Pedigree(t *testing.T) {
}

func TestBuild_Evidence(t *testing.T) {
t.Parallel()
api := &mockModelAPIFetcher{fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) {
return &fetcher.ModelAPIResponse{ID: id, SHA: "abc"}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

// Source scan: the same model found in two files.
d := scanner.Discovery{ID: "org/m", Name: "org/m", Type: "model", Occurrences: []scanner.Occurrence{
{Location: "a/one.py", Line: 3, Method: "from_pretrained", Symbol: "org/m"},
{Location: "two.yaml", Line: 1, Method: "yaml_model_field", Symbol: "org/m"},
}}
boms, err := BuildPerDiscovery([]scanner.Discovery{d}, GenerateOptions{})
boms, err := g.buildPerDiscovery([]scanner.Discovery{d}, GenerateOptions{})
if err != nil || len(boms) != 1 {
t.Fatalf("BuildPerDiscovery = %d boms, err %v", len(boms), err)
}
Expand All @@ -291,7 +298,7 @@ func TestBuild_Evidence(t *testing.T) {
}

// Model-ID input: identified through the Hub API.
boms, err = BuildFromModelIDs([]string{"org/m"}, GenerateOptions{})
boms, err = g.buildFromModelIDs([]string{"org/m"}, GenerateOptions{})
if err != nil || len(boms) != 1 {
t.Fatalf("BuildFromModelIDs = %d boms, err %v", len(boms), err)
}
Expand All @@ -311,12 +318,13 @@ func TestBuild_Evidence(t *testing.T) {
// A dataset known only from a dataset:<id> tag gets a component, and the card's
// datasets entry (built from the same tag) references it.
func TestBuildFromModelIDs_DatasetTagFallback(t *testing.T) {
t.Parallel()
api := &mockModelAPIFetcher{fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) {
return &fetcher.ModelAPIResponse{ID: "org/model", SHA: "abc", Tags: []string{"pytorch", "dataset:wikipedia"}}, nil
}}
withFetchers(t, hfLikeFetchers(api))
g := withFetchers(hfLikeFetchers(api))

boms, err := BuildFromModelIDs([]string{"org/model"}, GenerateOptions{})
boms, err := g.buildFromModelIDs([]string{"org/model"}, GenerateOptions{})
if err != nil || len(boms) != 1 {
t.Fatalf("BuildFromModelIDs = %d boms, err %v", len(boms), err)
}
Expand Down
Loading
Loading