l_ai_category/ai_test.go

135 lines
2.9 KiB
Go

package l_ai_category
import (
"context"
"fmt"
"gitea.cdlsxd.cn/self-tools/l_ai_category/utils_gorm"
"sync"
"testing"
"time"
)
func TestCategoryGoodsSet(t *testing.T) {
var wg sync.WaitGroup
defer clean()
path := getPath()
goods := goodsAll()
updateChan := make(chan *UpdateLog, len(goods))
start := time.Now().Unix()
wg.Add(len(goods))
for _, v := range goods {
go func(updateChan chan *UpdateLog) {
defer wg.Done()
res, err := GetCategory(context.Background(), v.Title, key, modelType, path)
if err != nil || res == 0 {
return
}
updateChan <- &UpdateLog{
GoodsId: v.Id,
Title: v.Title,
CateId: res,
}
}(updateChan)
}
wg.Wait()
close(updateChan)
end := time.Now().Unix()
fmt.Printf("耗时:%d秒", end-start)
updateSlice := make([]*UpdateLog, len(updateChan))
for v := range updateChan {
updateSlice = append(updateSlice, v)
}
for _, v := range updateSlice {
if v == nil {
continue
}
updateGoodsCate(v)
}
t.Log(updateSlice)
}
func updateGoodsCate(data *UpdateLog) {
db.Table("goods_category_relation").Where("goods_id = ?", data.GoodsId).Updates(map[string]interface{}{"category_id": data.CateId})
db.Table("goods").Where("id = ?", data.GoodsId).Updates(map[string]interface{}{"brand_category_check": 3})
}
type UpdateLog struct {
GoodsId int `json:"goods_id"`
Title string `json:"title"`
CateId int `json:"cate_id"`
}
type CategoryPhy struct {
Id int `json:"id"`
Name string `json:"name"`
Level int `json:"level"`
ParentId int `json:"parent_id"`
}
type Goods struct {
Id int `json:"id"`
Title string `json:"title"`
BrandCateGoryCheck int `json:"brand_category_check"`
}
func goodsAll() (data []Goods) {
var (
row Goods
err error
)
rows, _ := db.Raw("SELECT id,title FROM `goods` where brand_category_check=4").Rows()
defer rows.Close()
for rows.Next() {
err = db.ScanRows(rows, &row)
if err != nil {
panic(err)
}
data = append(data, row)
}
return
}
func getPath() []map[int][]*CategoryDic {
cates := cateAll()
var (
cateLevelPath = make([]map[int][]*CategoryDic, 3)
)
for _, v := range cates {
if cateLevelPath[v.Level-1] == nil {
cateLevelPath[v.Level-1] = make(map[int][]*CategoryDic)
}
if _, ex := cateLevelPath[v.Level-1][v.ParentId]; !ex {
cateLevelPath[v.Level-1][v.ParentId] = make([]*CategoryDic, 0)
}
cateLevelPath[v.Level-1][v.ParentId] = append(cateLevelPath[v.Level-1][v.ParentId], &CategoryDic{
Category{
Id: v.Id,
Name: v.Name,
Level: v.Level,
}})
}
return cateLevelPath
}
func cateAll() (data []CategoryPhy) {
var (
row CategoryPhy
)
rows, _ := db.Raw(fmt.Sprintf("SELECT id,name,level,parent_id FROM `%s`", table)).Rows()
defer rows.Close()
for rows.Next() {
err := db.ScanRows(rows, &row)
if err != nil {
panic(err)
}
data = append(data, row)
}
return
}
var db, _, clean = utils_gorm.DB(driver, dns)